【Bug已解决】Kosmos2.5： index error on long ocr input 解决方案-育师

【Bug已解决】Kosmos2.5： index error on long ocr input 解决方案

一、现象长什么样

Kosmos2.5 是一个面向 OCR/文档理解的多模态模型，输入通常是"一张文档图 + 一段提示文本"。当文档较长（高分辨率扫描、多行密集文字）时，generate或forward抛索引错误：

# 现象 A：图像 patch 索引越界 IndexError: index 1024 is out of bounds for dimension 0 with size 1024 File ".../models/kosmos2_5/modeling_kosmos2_5.py", line 210, in forward img_feat = image_features[image_token_indices] # 现象 B：截断后位置错位 RuntimeError: index -1 is out of bounds for dimension 0 with size 0 # 长输入被截断到 max_length，但 image_token 的占位索引还指着被截掉的位置 # 现象 C：batch 内长短不一拼 padding 后越界 ValueError: too many indices for tensor of dimension 1 # padding 把序列补齐到统一长度，但 image_token_indices 仍是原始未 padding 的下标

最典型的触发：一张高分辨率文档图被切成很多 patch（比如 1024 个），加上 OCR 提示文本后总长度超过max_length=2048，截断逻辑只截了文本侧，却忘了同步修正图像 token 的索引 → 索引越界。

二、背景

Kosmos2.5 的输入构造流程是：

图像经 vision encoder 切成若干 patch 特征（image_features，长度 = patch 数 N）。
文本里用特殊<image>token 占位， tokenizer 后这些占位被展开成 N 个 image token，分布在序列的不同位置。
forward 时，模型根据image_token_indices（这些 image token 在序列里的下标）去image_features里取对应特征，拼回序列。

这套机制依赖一个不变式：image_token_indices里的最大值 < N（patch 总数），且截断/ padding 后这些索引必须同步更新。当长 OCR 输入触发截断（为了塞进max_length）或 padding（为了 batch），这个不变式被打破，就出现上面的索引错误。

三、根因

根因有三类：

截断只截断文本，不修正 image_token_indices。processor在max_length超限时直接裁掉文本 token 尾部，但 image token 的下标是相对"裁剪前序列"算的。裁剪后序列变短，那些指向被裁区域的 image token 下标变成非法（指向越界或负位置）。
padding 后索引未平移/未 mask。 batch 推理时短样本被 pad 到最长。padding 在序列前面或后面插入了 dummy token，但image_token_indices仍是原始下标，没有随 padding 偏移 → 在 padding 之后的位置取 image_features 时错位。
patch 数 N 与 image token 数不一致。长文档切的 patch 数超过image_features实际长度（例如 image encoder 有自己内部的max_patches限制，超出的 patch 被丢弃），但 tokenizer 展开的 image token 数仍是按"未限制"算的 → 下标越界。

四、最小可运行复现

下面用纯 Python 模拟"截断文本后 image_token_indices 越界"的逻辑：

from typing import List def build_image_token_indices(seq_len: int, n_image_tokens: int) -> List[int]: """模拟：在序列末尾均匀放置 n_image_tokens 个 image token 的下标。""" step = max(1, seq_len // n_image_tokens) return list(range(0, seq_len, step))[:n_image_tokens] def truncate(seq_len: int, max_len: int) -> int: return min(seq_len, max_len) # 正常短输入 seq_len = 500 n_img = 100 idx = build_image_token_indices(seq_len, n_img) print("短输入 max idx:", max(idx), "patch 数 N:", n_img) # 合法 # 长输入触发截断 max_len = 300 new_len = truncate(seq_len, max_len) new_idx = build_image_token_indices(seq_len, n_img) # 索引仍按旧 seq_len 算！ print("截断后序列长:", new_len, "但 image idx max:", max(new_idx)) assert max(new_idx) >= new_len, "复现成功：截断后 image token 索引越界" # 修正版：截断时同步裁剪 image_token_indices def truncate_with_indices(seq_len, max_len, idx): return [i for i in idx if i < max_len] fixed = truncate_with_indices(seq_len, max_len, idx) print("修正后 image idx:", fixed, "max:", max(fixed) if fixed else None) assert all(i < new_len for i in fixed), "修正失败"

运行后，原new_idx的最大值（~499）超过了截断后的序列长度（300），触发越界；修正函数把越界索引裁掉，恢复不变式。

五、解决方案（第一层：最小直接修复）

最快的止血：在调用 processor / 截断前，手动清洗 image_token_indices，使其始终落在有效范围内：

import torch def sanitize_image_token_indices(image_token_indices, seq_len, image_features_len): """第一层修复：保证索引在 [0, seq_len) 且 < image_features_len。""" valid = [] for i in image_token_indices: if 0 <= i < seq_len and i < image_features_len: valid.append(i) # 若全部越界（极端长输入），退化为均匀取样 image_features if not valid and image_features_len > 0: step = max(1, image_features_len // seq_len) if seq_len else 1 valid = list(range(0, image_features_len, max(step, 1)))[:seq_len] return valid # 使用示意：在构造模型输入后、forward 前 inputs = processor(images=doc_image, text=prompt, return_tensors="pt", truncation=True, max_length=2048) seq_len = inputs["input_ids"].shape[1] image_token_indices = (inputs["input_ids"][0] == processor.image_token_id).nonzero().flatten().tolist() # 取出 image_features（来自 vision encoder） valid_idx = sanitize_image_token_indices(image_token_indices, seq_len, image_features.shape[0]) # 用合法索引重建（或传给模型一个已清洗的 indices 参数） assert max(valid_idx) < image_features.shape[0], "仍有越界，检查 image_features 长度"

第一层让用户立刻消除IndexError，长 OCR 文档也能跑。

六、解决方案（第二层：结构性改进）

把"索引与序列同步"做成KosmosIndexSync，在 processor 和模型之间统一维护不变式：

from dataclasses import dataclass from typing import List @dataclass class KosmosIndexSync: """维护 image_token_indices 与(截断后)序列长度、image_features 长度的一致。""" max_patches: int = 1024 def sync_after_truncation(self, indices: List[int], new_seq_len: int) -> List[int]: kept = [i for i in indices if 0 <= i < new_seq_len] # 同时保证不超过 image_features 实际容量 kept = [i for i in kept if i < self.max_patches] # 若因截断丢失过多 image token，从 image_features 均匀补回 if len(kept) < max(1, len(indices) // 2) and self.max_patches > 0: step = max(1, self.max_patches // max(new_seq_len, 1)) kept = list(range(0, self.max_patches, step))[:new_seq_len] return kept def sync_after_padding(self, indices: List[int], pad_left: int) -> List[int]: # padding 在左侧插入 dummy 时，所有索引右移 pad_left return [i + pad_left for i in indices] # 使用 sync = KosmosIndexSync(max_patches=1024) inputs = processor(images=doc_image, text=prompt, return_tensors="pt", truncation=True, max_length=2048, padding="max_length") seq_len = inputs["input_ids"].shape[1] raw_idx = (inputs["input_ids"][0] == processor.image_token_id).nonzero().flatten().tolist() valid = sync.sync_after_truncation(raw_idx, seq_len) if inputs.get("attention_mask") is not None: pad_left = int((inputs["attention_mask"][0] == 0).sum().item()) # 左 padding 数量 valid = sync.sync_after_padding(valid, pad_left)

KosmosIndexSync把"截断同步 + 左 padding 平移 + 容量上限"三件事集中处理，保证image_token_indices永远落在合法区间。

七、解决方案（第三层：断言 / CI 守护）

用 pytest 固化"长输入不产生越界索引"的契约：

import pytest def test_indices_within_bounds_after_truncation(): from index_sync import KosmosIndexSync sync = KosmosIndexSync(max_patches=1024) # 模拟长 OCR：1024 个 image token，序列被截到 300 indices = list(range(0, 10000, 10))[:1024] new_len = 300 valid = sync.sync_after_truncation(indices, new_len) assert all(0 <= i < new_len for i in valid), "截断后索引仍越界" assert all(i < 1024 for i in valid), "索引超过 image_features 容量" def test_padding_shifts_indices(): from index_sync import KosmosIndexSync sync = KosmosIndexSync() idx = [5, 10, 15] shifted = sync.sync_after_padding(idx, pad_left=4) assert shifted == [9, 14, 19], "左 padding 后索引应整体右移" def test_no_indexerror_on_long_ocr(): # 端到端：长文档不应抛 IndexError import torch from unittest.mock import MagicMock image_features = torch.randn(1024, 64, 64) # N=1024 patch indices = list(range(0, 10000, 10))[:1024] new_len = 300 valid = [i for i in indices if i < new_len and i < image_features.shape[0]] # 取特征不应越界 feats = image_features[valid] assert feats.shape[0] == len(valid)

CI 跑pytest tests/test_kosmos2_5_long_ocr.py，以后只要截断/padding 逻辑又忘了同步索引，测试立刻红灯。

八、排查清单

当 Kosmos2.5 在长 OCR 输入上报索引错误，按顺序查：

IndexError: index X is out of bounds for image_features→image_token_indices越界，先用sanitize_image_token_indices清洗。
index -1 is out of bounds→ 截断把 image token 全裁掉了，需要同步裁剪或均匀补回。
batch 推理报too many indices→ padding 后索引未平移，用sync_after_padding右移。
确认image_features实际 patch 数 N，与 tokenizer 展开的 image token 数是否一致；不一致要限制max_patches。
长期方案：把索引同步逻辑收进 processor（返回已清洗的 indices），而不是让模型侧去猜。