Skip to content

Commit ea82cff

Browse files
author
zhansheng.lzs
committed
Fix tiktoken StackOverflow on long text inputs
tiktoken's recursive BPE merging overflows the call stack on very long inputs, causing pyo3_runtime.PanicException: StackOverflow. Split text into chunks (<= 100k chars) at line boundaries before encoding. Short text takes the original fast path unchanged. Fixes #116
1 parent eadf165 commit ea82cff

1 file changed

Lines changed: 42 additions & 5 deletions

File tree

dashscope/tokenizers/qwen_tokenizer.py

Lines changed: 42 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,11 @@
3333
)
3434
SPECIAL_TOKENS_SET = set(t for i, t in SPECIAL_TOKENS)
3535

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+
3641

3742
class QwenTokenizer(Tokenizer):
3843
@staticmethod
@@ -102,11 +107,43 @@ def encode( # type: ignore[override]
102107
disallowed_special: Union[Collection, str] = (),
103108
) -> Union[List[List], List]:
104109
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
110147

111148
def decode(
112149
self,

0 commit comments

Comments
 (0)