|
33 | 33 | ) |
34 | 34 | SPECIAL_TOKENS_SET = set(t for i, t in SPECIAL_TOKENS) |
35 | 35 |
|
| 36 | +# tiktoken's BPE merges tokens recursively in Rust, which can overflow the |
| 37 | +# call stack on very long inputs (pyo3_runtime.PanicException: StackOverflow). |
| 38 | +# Split text into chunks below this threshold before encoding. |
| 39 | +_CHUNK_SIZE = 100_000 |
| 40 | + |
36 | 41 |
|
37 | 42 | class QwenTokenizer(Tokenizer): |
38 | 43 | @staticmethod |
@@ -102,11 +107,43 @@ def encode( # type: ignore[override] |
102 | 107 | disallowed_special: Union[Collection, str] = (), |
103 | 108 | ) -> Union[List[List], List]: |
104 | 109 | text = unicodedata.normalize("NFC", text) |
105 | | - return self._tokenizer.encode( |
106 | | - text, |
107 | | - allowed_special=allowed_special, |
108 | | - disallowed_special=disallowed_special, |
109 | | - ) |
| 110 | + if len(text) <= _CHUNK_SIZE: |
| 111 | + return self._tokenizer.encode( |
| 112 | + text, |
| 113 | + allowed_special=allowed_special, |
| 114 | + disallowed_special=disallowed_special, |
| 115 | + ) |
| 116 | + |
| 117 | + result = [] |
| 118 | + for chunk in self._split_text(text): |
| 119 | + result.extend( |
| 120 | + self._tokenizer.encode( |
| 121 | + chunk, |
| 122 | + allowed_special=allowed_special, |
| 123 | + disallowed_special=disallowed_special, |
| 124 | + ), |
| 125 | + ) |
| 126 | + return result |
| 127 | + |
| 128 | + @staticmethod |
| 129 | + def _split_text(text: str, chunk_size: int = _CHUNK_SIZE) -> List[str]: |
| 130 | + """Split text into chunks at safe tokenization boundaries.""" |
| 131 | + parts: List[str] = [] |
| 132 | + for i, line in enumerate(text.split("\n")): |
| 133 | + piece = line if i == 0 else "\n" + line |
| 134 | + if len(piece) <= chunk_size: |
| 135 | + parts.append(piece) |
| 136 | + else: |
| 137 | + for j in range(0, len(piece), chunk_size): |
| 138 | + parts.append(piece[j : j + chunk_size]) |
| 139 | + |
| 140 | + chunks: List[str] = [] |
| 141 | + for part in parts: |
| 142 | + if chunks and len(chunks[-1]) + len(part) <= chunk_size: |
| 143 | + chunks[-1] += part |
| 144 | + else: |
| 145 | + chunks.append(part) |
| 146 | + return chunks |
110 | 147 |
|
111 | 148 | def decode( |
112 | 149 | self, |
|
0 commit comments