Size vocab by max token id so duplicate-piece vocabs (Yi-6B) don't overflow embedding

2026-08-26 14:42:12 +08:00
parent 94d0f2ff6a
commit ea7167b3f7
1 changed files with 5 additions and 1 deletions
@@ -57,7 +57,11 @@ class HuggingFaceTokenizer:

    @property
    def vocab_size(self) -> int:
-        return int(len(self._tok))
+        # len(tok) 数的是去重后的 surface form; 词表有重复 piece 时 (如 Yi-6B
+        # 63992 vs 最大 id 63999) 会小于真实 id 范围, embedding 越界触发
+        # device-side assert. 以最大 id + 1 为准.
+        max_id = max(self._tok.get_vocab().values())
+        return max(int(len(self._tok)), max_id + 1)

    def encode(self, text: str) -> list[int]:
        return list(self._tok.encode(text, add_special_tokens=False))