| # Copyright 2026 The IREE Authors | |
| # | |
| # Licensed under the Apache License v2.0 with LLVM Exceptions. | |
| # See https://llvm.org/LICENSE.txt for license information. | |
| # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | |
| import base64 | |
| import functools | |
| import gc | |
| import json | |
| import logging | |
| import os | |
| import tempfile | |
| import threading | |
| import unittest | |
| from pathlib import Path | |
| import iree.runtime as rt | |
| TESTDATA_DIR = ( | |
| Path(__file__).resolve().parent.parent.parent.parent.parent | |
| / "runtime" | |
| / "src" | |
| / "iree" | |
| / "tokenizer" | |
| / "testdata" | |
| ) | |
| BPE_MINIMAL_JSON = TESTDATA_DIR / "bpe_bytelevel_minimal.json" | |
| BPE_UNICODE_JSON = TESTDATA_DIR / "bpe_bytelevel_unicode.json" | |
| def _require_testdata(test_func): | |
| """Skip test if tokenizer testdata is not available.""" | |
| @functools.wraps(test_func) | |
| def wrapper(self): | |
| if not BPE_MINIMAL_JSON.exists(): | |
| self.skipTest(f"Tokenizer testdata not found: {BPE_MINIMAL_JSON}") | |
| return test_func(self) | |
| return wrapper | |
| class TokenizerLoadTest(unittest.TestCase): | |
| @_require_testdata | |
| def test_from_file(self): | |
| tok = rt.Tokenizer.from_file(str(BPE_MINIMAL_JSON)) | |
| self.assertEqual(tok.model_type, "BPE") | |
| self.assertGreater(tok.vocab_size, 0) | |
| @_require_testdata | |
| def test_from_file_pathlike(self): | |
| """from_file() should accept Path objects, not just str.""" | |
| tok = rt.Tokenizer.from_file(BPE_MINIMAL_JSON) # Path, not str | |
| self.assertEqual(tok.model_type, "BPE") | |
| @_require_testdata | |
| def test_from_huggingface_json(self): | |
| json_str = BPE_MINIMAL_JSON.read_text() | |
| tok = rt.Tokenizer.from_huggingface_json(json_str) | |
| self.assertEqual(tok.model_type, "BPE") | |
| def test_from_file_not_found(self): | |
| with self.assertRaises(ValueError): | |
| rt.Tokenizer.from_file("/nonexistent/path/tokenizer.json") | |
| def test_from_file_bad_format(self): | |
| with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: | |
| f.write("not json content") | |
| path = f.name | |
| try: | |
| with self.assertRaises(ValueError): | |
| rt.Tokenizer.from_file(path) | |
| finally: | |
| Path(path).unlink() | |
| def test_from_huggingface_json_invalid(self): | |
| with self.assertRaises(Exception): | |
| rt.Tokenizer.from_huggingface_json("{invalid json") | |
| @_require_testdata | |
| def test_from_file_bytes_path(self): | |
| """from_file accepts bytes paths via os.fsencode.""" | |
| tok = rt.Tokenizer.from_file(os.fsencode(str(BPE_MINIMAL_JSON))) | |
| self.assertEqual(tok.model_type, "BPE") | |
| def test_from_file_rejects_non_path(self): | |
| """from_file raises TypeError for non-path types.""" | |
| with self.assertRaises(TypeError): | |
| rt.Tokenizer.from_file(123) | |
| with self.assertRaises(TypeError): | |
| rt.Tokenizer.from_file(object()) | |
| def test_from_tiktoken_inline(self): | |
| """from_tiktoken creates a tokenizer from inline tiktoken data.""" | |
| # Tiktoken requires all 256 single-byte tokens as the base vocabulary. | |
| lines = [] | |
| for byte_val in range(256): | |
| b64 = base64.b64encode(bytes([byte_val])).decode() | |
| lines.append(f"{b64} {byte_val}") | |
| data = "\n".join(lines) | |
| tok = rt.Tokenizer.from_tiktoken(data, "cl100k_base") | |
| self.assertGreaterEqual(tok.vocab_size, 256) | |
| def test_from_file_tiktoken_unknown_encoding(self): | |
| """from_file with .tiktoken extension and unknown encoding raises.""" | |
| with tempfile.NamedTemporaryFile( | |
| mode="w", suffix=".tiktoken", prefix="unknown_enc", delete=False | |
| ) as f: | |
| f.write("AA== 0\n") | |
| path = f.name | |
| try: | |
| with self.assertRaises(ValueError): | |
| rt.Tokenizer.from_file(path) | |
| finally: | |
| Path(path).unlink() | |
| @_require_testdata | |
| def test_multiple_tokenizers(self): | |
| """Verify independent tokenizer instances work correctly.""" | |
| tok1 = rt.Tokenizer.from_file(str(BPE_MINIMAL_JSON)) | |
| tok2 = rt.Tokenizer.from_file(str(BPE_MINIMAL_JSON)) | |
| ids1 = tok1.encode("Hello") | |
| ids2 = tok2.encode("Hello") | |
| self.assertEqual(ids1, ids2) | |
| del tok1 | |
| gc.collect() | |
| # tok2 should still work after tok1 is destroyed. | |
| ids3 = tok2.encode("Hello") | |
| self.assertEqual(ids2, ids3) | |
| class TokenizerEncodeDecodeTest(unittest.TestCase): | |
| @classmethod | |
| def setUpClass(cls): | |
| if not BPE_MINIMAL_JSON.exists(): | |
| raise unittest.SkipTest(f"Tokenizer testdata not found: {BPE_MINIMAL_JSON}") | |
| cls.tok = rt.Tokenizer.from_file(str(BPE_MINIMAL_JSON)) | |
| def test_encode_decode_roundtrip(self): | |
| text = "Hello world" | |
| ids = self.tok.encode(text) | |
| self.assertIsInstance(ids, list) | |
| self.assertTrue(all(isinstance(i, int) for i in ids)) | |
| decoded = self.tok.decode(ids) | |
| self.assertEqual(decoded, text) | |
| def test_encode_empty(self): | |
| ids = self.tok.encode("") | |
| self.assertEqual(ids, []) | |
| def test_decode_empty(self): | |
| text = self.tok.decode([]) | |
| self.assertEqual(text, "") | |
| def test_encode_multiple_words(self): | |
| text = "The quick brown fox" | |
| ids = self.tok.encode(text) | |
| decoded = self.tok.decode(ids) | |
| self.assertEqual(decoded, text) | |
| def test_encode_add_special_tokens_no_postprocessor(self): | |
| # Minimal tokenizer has no post-processor, so add_special_tokens | |
| # should produce the same output (no BOS/EOS to add). | |
| ids_without = self.tok.encode("Hello") | |
| ids_with = self.tok.encode("Hello", add_special_tokens=True) | |
| self.assertEqual(ids_without, ids_with) | |
| def test_decode_skip_special_tokens(self): | |
| ids = self.tok.encode("Hello") | |
| text_without = self.tok.decode(ids) | |
| text_with = self.tok.decode(ids, skip_special_tokens=True) | |
| self.assertEqual(text_without, text_with) | |
| @unittest.skip( | |
| "ByteLevel decoder STATELESS bug: decode drops multi-byte UTF-8. " | |
| "Fix: change CAPABILITY_STATELESS to CAPABILITY_NONE in byte_level.c." | |
| ) | |
| def test_encode_decode_unicode(self): | |
| """Byte-level BPE must round-trip non-ASCII text (multibyte UTF-8).""" | |
| if not BPE_UNICODE_JSON.exists(): | |
| self.skipTest(f"Unicode tokenizer not found: {BPE_UNICODE_JSON}") | |
| tok = rt.Tokenizer.from_file(str(BPE_UNICODE_JSON)) | |
| for text in ["café", "你好世界", "hello 😀 world", "Ñoño"]: | |
| with self.subTest(text=text): | |
| ids = tok.encode(text) | |
| self.assertTrue(len(ids) > 0) | |
| self.assertEqual(tok.decode(ids), text) | |
| def test_repr(self): | |
| r = repr(self.tok) | |
| self.assertIn("Tokenizer", r) | |
| self.assertIn("BPE", r) | |
| self.assertIn("112", r) | |
| class TokenizerStreamingEncodeTest(unittest.TestCase): | |
| @classmethod | |
| def setUpClass(cls): | |
| if not BPE_MINIMAL_JSON.exists(): | |
| raise unittest.SkipTest(f"Tokenizer testdata not found: {BPE_MINIMAL_JSON}") | |
| cls.tok = rt.Tokenizer.from_file(str(BPE_MINIMAL_JSON)) | |
| def test_streaming_encode_matches_batch(self): | |
| batch_ids = self.tok.encode("Hello world") | |
| enc = self.tok.encode_stream() | |
| ids1 = enc.feed("Hello ") | |
| ids2 = enc.feed("world") | |
| ids3 = enc.finalize() | |
| self.assertEqual(ids1 + ids2 + ids3, batch_ids) | |
| def test_context_manager(self): | |
| with self.tok.encode_stream() as enc: | |
| ids = enc.feed("test") | |
| ids += enc.finalize() | |
| self.assertIsInstance(ids, list) | |
| self.assertGreater(len(ids), 0) | |
| def test_finalize_twice_raises(self): | |
| enc = self.tok.encode_stream() | |
| enc.feed("x") | |
| enc.finalize() | |
| with self.assertRaises(ValueError): | |
| enc.finalize() | |
| def test_feed_after_finalize_raises(self): | |
| enc = self.tok.encode_stream() | |
| enc.finalize() | |
| with self.assertRaises(ValueError): | |
| enc.feed("x") | |
| def test_multiple_streams_sequential(self): | |
| """Verify creating multiple streams from the same tokenizer works.""" | |
| for _ in range(3): | |
| enc = self.tok.encode_stream() | |
| ids = enc.feed("test") | |
| ids += enc.finalize() | |
| del enc | |
| gc.collect() | |
| # Tokenizer should still work after all streams are destroyed. | |
| self.assertGreater(len(self.tok.encode("test")), 0) | |
| class TokenizerStreamingDecodeTest(unittest.TestCase): | |
| @classmethod | |
| def setUpClass(cls): | |
| if not BPE_MINIMAL_JSON.exists(): | |
| raise unittest.SkipTest(f"Tokenizer testdata not found: {BPE_MINIMAL_JSON}") | |
| cls.tok = rt.Tokenizer.from_file(str(BPE_MINIMAL_JSON)) | |
| cls.ids = cls.tok.encode("Hello world") | |
| def test_streaming_decode_matches_batch(self): | |
| dec = self.tok.decode_stream() | |
| t1 = dec.feed(self.ids[:2]) | |
| t2 = dec.feed(self.ids[2:]) | |
| t3 = dec.finalize() | |
| self.assertEqual(t1 + t2 + t3, "Hello world") | |
| def test_context_manager(self): | |
| with self.tok.decode_stream() as dec: | |
| text = dec.feed(self.ids) | |
| text += dec.finalize() | |
| self.assertEqual(text, "Hello world") | |
| def test_finalize_twice_raises(self): | |
| dec = self.tok.decode_stream() | |
| dec.feed(self.ids) | |
| dec.finalize() | |
| with self.assertRaises(ValueError): | |
| dec.finalize() | |
| def test_feed_after_finalize_raises(self): | |
| dec = self.tok.decode_stream() | |
| dec.finalize() | |
| with self.assertRaises(ValueError): | |
| dec.feed(self.ids) | |
| def test_single_token_feed(self): | |
| """Feed tokens one at a time and verify accumulated result.""" | |
| dec = self.tok.decode_stream() | |
| parts = [] | |
| for token_id in self.ids: | |
| parts.append(dec.feed([token_id])) | |
| parts.append(dec.finalize()) | |
| self.assertEqual("".join(parts), "Hello world") | |
| def test_feed_one(self): | |
| """feed_one() should produce the same result as feed([id]).""" | |
| dec_list = self.tok.decode_stream() | |
| dec_one = self.tok.decode_stream() | |
| parts_list = [] | |
| parts_one = [] | |
| for token_id in self.ids: | |
| parts_list.append(dec_list.feed([token_id])) | |
| parts_one.append(dec_one.feed_one(token_id)) | |
| parts_list.append(dec_list.finalize()) | |
| parts_one.append(dec_one.finalize()) | |
| self.assertEqual("".join(parts_list), "".join(parts_one)) | |
| self.assertEqual("".join(parts_one), "Hello world") | |
| def test_feed_one_after_finalize_raises(self): | |
| dec = self.tok.decode_stream() | |
| dec.finalize() | |
| with self.assertRaises(ValueError): | |
| dec.feed_one(0) | |
| class TokenizerVocabTest(unittest.TestCase): | |
| @classmethod | |
| def setUpClass(cls): | |
| if not BPE_MINIMAL_JSON.exists(): | |
| raise unittest.SkipTest(f"Tokenizer testdata not found: {BPE_MINIMAL_JSON}") | |
| cls.tok = rt.Tokenizer.from_file(str(BPE_MINIMAL_JSON)) | |
| def test_vocab_size(self): | |
| self.assertEqual(self.tok.vocab_size, 112) | |
| def test_model_type(self): | |
| self.assertEqual(self.tok.model_type, "BPE") | |
| def test_id_to_token_valid(self): | |
| token = self.tok.id_to_token(39) | |
| self.assertIsNotNone(token) | |
| self.assertIsInstance(token, str) | |
| def test_id_to_token_out_of_range(self): | |
| self.assertIsNone(self.tok.id_to_token(99999)) | |
| def test_id_to_token_negative(self): | |
| self.assertIsNone(self.tok.id_to_token(-1)) | |
| self.assertIsNone(self.tok.id_to_token(-100)) | |
| def test_token_roundtrip(self): | |
| """Verify id_to_token and encode are consistent.""" | |
| ids = self.tok.encode("H") | |
| self.assertEqual(len(ids), 1) | |
| token_text = self.tok.id_to_token(ids[0]) | |
| self.assertEqual(token_text, "H") | |
| def test_token_to_id_known(self): | |
| """token_to_id returns correct ID for a known vocab entry.""" | |
| ids = self.tok.encode("H") | |
| self.assertEqual(len(ids), 1) | |
| looked_up = self.tok.token_to_id("H") | |
| self.assertIsNotNone(looked_up) | |
| self.assertEqual(looked_up, ids[0]) | |
| def test_token_to_id_unknown(self): | |
| self.assertIsNone(self.tok.token_to_id("nonexistent_token_xyz")) | |
| def test_special_ids(self): | |
| ids = self.tok.special_ids | |
| self.assertIsInstance(ids, dict) | |
| for key in ("bos", "eos", "unk", "pad", "sep", "cls", "mask"): | |
| self.assertIn(key, ids) | |
| # Value is int (if configured) or None (if absent). | |
| self.assertTrue( | |
| ids[key] is None or isinstance(ids[key], int), | |
| f"special_ids[{key!r}] = {ids[key]!r}, expected int or None", | |
| ) | |
| class TokenizerThreadingTest(unittest.TestCase): | |
| """Test concurrent use of a shared Tokenizer from multiple threads. | |
| The Tokenizer itself is thread-safe (immutable C object). Each thread | |
| must use its own EncodeStream/DecodeStream (not shared across threads). | |
| """ | |
| @classmethod | |
| def setUpClass(cls): | |
| if not BPE_MINIMAL_JSON.exists(): | |
| raise unittest.SkipTest(f"Tokenizer testdata not found: {BPE_MINIMAL_JSON}") | |
| cls.tok = rt.Tokenizer.from_file(str(BPE_MINIMAL_JSON)) | |
| def test_concurrent_batch_encode(self): | |
| import threading | |
| results = {} | |
| errors = [] | |
| def worker(thread_id): | |
| try: | |
| text = f"Thread {thread_id} hello world" | |
| for _ in range(100): | |
| ids = self.tok.encode(text) | |
| results[thread_id] = ids | |
| except Exception as e: | |
| errors.append((thread_id, e)) | |
| threads = [threading.Thread(target=worker, args=(i,)) for i in range(4)] | |
| for t in threads: | |
| t.start() | |
| for t in threads: | |
| t.join() | |
| self.assertEqual(len(errors), 0, f"Thread errors: {errors}") | |
| for i in range(4): | |
| expected = self.tok.encode(f"Thread {i} hello world") | |
| self.assertEqual(results[i], expected) | |
| def test_concurrent_streaming(self): | |
| import threading | |
| results = {} | |
| errors = [] | |
| def worker(thread_id): | |
| try: | |
| text = f"Thread {thread_id} hello" | |
| enc = self.tok.encode_stream() | |
| ids = enc.feed(text) | |
| ids += enc.finalize() | |
| results[thread_id] = ids | |
| except Exception as e: | |
| errors.append((thread_id, e)) | |
| threads = [threading.Thread(target=worker, args=(i,)) for i in range(4)] | |
| for t in threads: | |
| t.start() | |
| for t in threads: | |
| t.join() | |
| self.assertEqual(len(errors), 0, f"Thread errors: {errors}") | |
| for i in range(4): | |
| expected = self.tok.encode(f"Thread {i} hello") | |
| self.assertEqual(results[i], expected) | |
| class TokenizerLargeInputTest(unittest.TestCase): | |
| """Test with inputs that exceed internal buffer sizes.""" | |
| @classmethod | |
| def setUpClass(cls): | |
| if not BPE_MINIMAL_JSON.exists(): | |
| raise unittest.SkipTest(f"Tokenizer testdata not found: {BPE_MINIMAL_JSON}") | |
| cls.tok = rt.Tokenizer.from_file(str(BPE_MINIMAL_JSON)) | |
| def test_streaming_encode_exceeds_token_buffer(self): | |
| """Feed a chunk that produces more tokens than the 256-token buffer.""" | |
| for n in [300, 500, 1000, 5000]: | |
| text = "a" * n | |
| batch_ids = self.tok.encode(text) | |
| enc = self.tok.encode_stream() | |
| stream_ids = enc.feed(text) | |
| stream_ids += enc.finalize() | |
| self.assertEqual( | |
| stream_ids, | |
| batch_ids, | |
| f"Stream/batch mismatch at {n} chars: " | |
| f"stream={len(stream_ids)}, batch={len(batch_ids)}", | |
| ) | |
| def test_large_multi_chunk_roundtrip(self): | |
| """120KB input split into 1KB chunks must match batch.""" | |
| text = "Hello world " * 10000 | |
| batch_ids = self.tok.encode(text) | |
| enc = self.tok.encode_stream() | |
| stream_ids = [] | |
| for i in range(0, len(text), 1024): | |
| stream_ids.extend(enc.feed(text[i : i + 1024])) | |
| stream_ids.extend(enc.finalize()) | |
| self.assertEqual(len(stream_ids), len(batch_ids)) | |
| self.assertEqual(stream_ids, batch_ids) | |
| def _build_special_tokens_json(): | |
| """Build a tokenizer JSON with BOS/EOS special tokens and a post-processor.""" | |
| base_json = BPE_MINIMAL_JSON.read_text() | |
| tok_dict = json.loads(base_json) | |
| # Add BOS (id=111 reused as <|bos|>) and EOS (id=112 as <|eos|>) | |
| tok_dict["model"]["vocab"]["<|bos|>"] = 112 | |
| tok_dict["model"]["vocab"]["<|eos|>"] = 113 | |
| tok_dict["added_tokens"].extend( | |
| [ | |
| { | |
| "id": 112, | |
| "content": "<|bos|>", | |
| "single_word": False, | |
| "lstrip": False, | |
| "rstrip": False, | |
| "normalized": False, | |
| "special": True, | |
| }, | |
| { | |
| "id": 113, | |
| "content": "<|eos|>", | |
| "single_word": False, | |
| "lstrip": False, | |
| "rstrip": False, | |
| "normalized": False, | |
| "special": True, | |
| }, | |
| ] | |
| ) | |
| tok_dict["post_processor"] = { | |
| "type": "TemplateProcessing", | |
| "single": [ | |
| {"SpecialToken": {"id": "<|bos|>", "type_id": 0}}, | |
| {"Sequence": {"id": "A", "type_id": 0}}, | |
| {"SpecialToken": {"id": "<|eos|>", "type_id": 0}}, | |
| ], | |
| "pair": [ | |
| {"SpecialToken": {"id": "<|bos|>", "type_id": 0}}, | |
| {"Sequence": {"id": "A", "type_id": 0}}, | |
| {"SpecialToken": {"id": "<|eos|>", "type_id": 0}}, | |
| {"Sequence": {"id": "B", "type_id": 1}}, | |
| {"SpecialToken": {"id": "<|eos|>", "type_id": 1}}, | |
| ], | |
| "special_tokens": { | |
| "<|bos|>": {"id": "<|bos|>", "ids": [112], "tokens": ["<|bos|>"]}, | |
| "<|eos|>": {"id": "<|eos|>", "ids": [113], "tokens": ["<|eos|>"]}, | |
| }, | |
| } | |
| return json.dumps(tok_dict) | |
| class TokenizerSpecialTokensTest(unittest.TestCase): | |
| """Test special token handling with a tokenizer that has BOS/EOS configured.""" | |
| @classmethod | |
| def setUpClass(cls): | |
| if not BPE_MINIMAL_JSON.exists(): | |
| raise unittest.SkipTest(f"Tokenizer testdata not found: {BPE_MINIMAL_JSON}") | |
| try: | |
| cls.tok = rt.Tokenizer.from_huggingface_json(_build_special_tokens_json()) | |
| except Exception as e: | |
| raise unittest.SkipTest(f"Failed to build special tokens tokenizer: {e}") | |
| def test_vocab_includes_special_tokens(self): | |
| # The added special tokens should be in the vocab. | |
| self.assertIsNotNone(self.tok.id_to_token(112)) # <|bos|> | |
| self.assertIsNotNone(self.tok.id_to_token(113)) # <|eos|> | |
| def test_add_special_tokens_adds_bos_eos(self): | |
| ids_without = self.tok.encode("Hello") | |
| ids_with = self.tok.encode("Hello", add_special_tokens=True) | |
| # With special tokens, the output should be longer (BOS + tokens + EOS). | |
| self.assertGreater(len(ids_with), len(ids_without)) | |
| def test_decode_skip_special_tokens(self): | |
| ids = self.tok.encode("Hello", add_special_tokens=True) | |
| text_with = self.tok.decode(ids, skip_special_tokens=False) | |
| text_skip = self.tok.decode(ids, skip_special_tokens=True) | |
| # Skipping special tokens should give clean text. | |
| self.assertIn("Hello", text_skip) | |
| # Not skipping should include special token text. | |
| self.assertGreater(len(text_with), len(text_skip)) | |
| def test_encode_special_token_in_text(self): | |
| """Special tokens in input text should be recognized.""" | |
| ids = self.tok.encode("<|endoftext|>") | |
| # Should produce the special token ID (111), not individual chars. | |
| self.assertIn(111, ids) | |
| def test_streaming_decode_skip_special_tokens(self): | |
| """Streaming decode with skip_special_tokens matches batch.""" | |
| ids = self.tok.encode("Hello", add_special_tokens=True) | |
| batch_skip = self.tok.decode(ids, skip_special_tokens=True) | |
| batch_noskip = self.tok.decode(ids, skip_special_tokens=False) | |
| # Sanity: batch skip should be shorter. | |
| self.assertGreater(len(batch_noskip), len(batch_skip)) | |
| # Streaming with skip. | |
| dec = self.tok.decode_stream(skip_special_tokens=True) | |
| stream_text = dec.feed(ids) + dec.finalize() | |
| self.assertEqual(stream_text, batch_skip) | |
| # Streaming without skip. | |
| dec2 = self.tok.decode_stream(skip_special_tokens=False) | |
| stream_noskip = dec2.feed(ids) + dec2.finalize() | |
| self.assertEqual(stream_noskip, batch_noskip) | |
| def test_encode_no_special_token_matching(self): | |
| """no_special_token_matching treats special tokens as literal text.""" | |
| ids_normal = self.tok.encode("<|endoftext|>") | |
| ids_ordinary = self.tok.encode("<|endoftext|>", no_special_token_matching=True) | |
| # Normal: matched as special token (111). | |
| self.assertIn(111, ids_normal) | |
| # Ordinary: NOT matched, tokenized as literal characters. | |
| self.assertNotIn(111, ids_ordinary) | |
| self.assertGreater(len(ids_ordinary), len(ids_normal)) | |
| def test_streaming_encode_add_special_tokens(self): | |
| """Streaming encode with add_special_tokens matches batch.""" | |
| text = "Hello" | |
| batch = self.tok.encode(text, add_special_tokens=True) | |
| enc = self.tok.encode_stream(add_special_tokens=True) | |
| stream = enc.feed(text) + enc.finalize() | |
| self.assertEqual(batch, stream) | |
| # Should have BOS (112) and EOS (113). | |
| self.assertIn(112, stream) | |
| self.assertIn(113, stream) | |
| def test_streaming_encode_no_special_token_matching(self): | |
| """Streaming encode with no_special_token_matching matches batch.""" | |
| text = "Hello <|endoftext|> world" | |
| batch = self.tok.encode(text, no_special_token_matching=True) | |
| enc = self.tok.encode_stream(no_special_token_matching=True) | |
| stream = enc.feed(text) + enc.finalize() | |
| self.assertEqual(batch, stream) | |
| self.assertNotIn(111, stream) | |
| class TokenizerPendingTokenBoundTest(unittest.TestCase): | |
| """Test pending_token_bound behavior via streaming encode edge cases. | |
| The C function pending_token_bound() is used internally by EncodeStream | |
| to size the finalize output buffer. These tests exercise edge cases | |
| where an incorrect bound would cause RESOURCE_EXHAUSTED errors. | |
| """ | |
| @classmethod | |
| def setUpClass(cls): | |
| if not BPE_MINIMAL_JSON.exists(): | |
| raise unittest.SkipTest(f"Tokenizer testdata not found: {BPE_MINIMAL_JSON}") | |
| cls.tok = rt.Tokenizer.from_file(str(BPE_MINIMAL_JSON)) | |
| def test_streaming_encode_large_input_finalize(self): | |
| """Feed 100KB+ of text in chunks, verify finalize succeeds and | |
| the round-trip matches batch encode. | |
| This tests that the tight bound from pending_token_bound is sufficient | |
| for large streams where the pipeline accumulates significant state. | |
| """ | |
| # Build a large text (>100KB). Use varied content to exercise different | |
| # BPE merge patterns and segmentation boundaries. | |
| base = "Hello world " | |
| repeat = (100 * 1024 // len(base)) + 1 | |
| text = base * repeat # ~100KB+ | |
| self.assertGreater(len(text), 100 * 1024) | |
| # Batch encode for ground truth. | |
| batch_ids = self.tok.encode(text) | |
| # Streaming encode in 1KB chunks. | |
| enc = self.tok.encode_stream() | |
| stream_ids = [] | |
| chunk_size = 1024 | |
| for i in range(0, len(text), chunk_size): | |
| stream_ids.extend(enc.feed(text[i : i + chunk_size])) | |
| # finalize() internally uses pending_token_bound to size its buffer. | |
| # If the bound is too small, this raises RESOURCE_EXHAUSTED. | |
| stream_ids.extend(enc.finalize()) | |
| self.assertEqual(len(stream_ids), len(batch_ids)) | |
| self.assertEqual(stream_ids, batch_ids) | |
| def test_streaming_encode_empty_finalize(self): | |
| """Create stream, call finalize immediately without any feed. | |
| Verify it succeeds and returns empty (no special tokens configured). | |
| """ | |
| enc = self.tok.encode_stream() | |
| result = enc.finalize() | |
| self.assertEqual(result, []) | |
| def test_streaming_encode_empty_finalize_with_special_tokens(self): | |
| """Create stream with add_special_tokens, call finalize immediately. | |
| Should succeed and return only the special tokens (BOS + EOS). | |
| """ | |
| try: | |
| tok = rt.Tokenizer.from_huggingface_json(_build_special_tokens_json()) | |
| except Exception as e: | |
| self.skipTest(f"Failed to build special tokens tokenizer: {e}") | |
| enc = tok.encode_stream(add_special_tokens=True) | |
| result = enc.finalize() | |
| # With BOS/EOS postprocessor and empty input, should get just BOS + EOS. | |
| self.assertIn(112, result) # BOS | |
| self.assertIn(113, result) # EOS | |
| def test_streaming_encode_single_byte_chunks(self): | |
| """Feed text one byte at a time — maximizes pipeline fragmentation. | |
| Exercises partial special token matches spanning many chunks. | |
| """ | |
| text = "Hello world, how are you today?" | |
| batch_ids = self.tok.encode(text) | |
| enc = self.tok.encode_stream() | |
| stream_ids = [] | |
| for ch in text: | |
| stream_ids.extend(enc.feed(ch)) | |
| stream_ids.extend(enc.finalize()) | |
| self.assertEqual(stream_ids, batch_ids) | |
| def test_streaming_encode_varied_chunk_sizes(self): | |
| """Feed text in irregular chunk sizes (1, 7, 3, 13, ...). | |
| Exercises pipeline boundary conditions at unpredictable positions. | |
| """ | |
| text = "The quick brown fox jumps over the lazy dog. " * 100 | |
| batch_ids = self.tok.encode(text) | |
| enc = self.tok.encode_stream() | |
| stream_ids = [] | |
| sizes = [1, 7, 3, 13, 2, 11, 5, 17, 4, 19] | |
| pos = 0 | |
| for i in range(len(text)): | |
| chunk_size = sizes[i % len(sizes)] | |
| chunk = text[pos : pos + chunk_size] | |
| if not chunk: | |
| break | |
| stream_ids.extend(enc.feed(chunk)) | |
| pos += chunk_size | |
| if pos >= len(text): | |
| break | |
| stream_ids.extend(enc.finalize()) | |
| self.assertEqual(stream_ids, batch_ids) | |
| def test_streaming_encode_special_tokens_in_chunks(self): | |
| """Feed text containing special tokens split across chunk boundaries. | |
| Exercises partial special token match recovery during finalize. | |
| """ | |
| try: | |
| tok = rt.Tokenizer.from_huggingface_json(_build_special_tokens_json()) | |
| except Exception as e: | |
| self.skipTest(f"Failed to build special tokens tokenizer: {e}") | |
| text = "Hello <|endoftext|> world" | |
| batch_ids = tok.encode(text) | |
| # Split right in the middle of <|endoftext|> | |
| for split_pos in range(len(text)): | |
| enc = tok.encode_stream() | |
| ids = enc.feed(text[:split_pos]) + enc.feed(text[split_pos:]) | |
| ids += enc.finalize() | |
| self.assertEqual(ids, batch_ids, f"Mismatch at split_pos={split_pos}") | |
| if __name__ == "__main__": | |
| logging.basicConfig(level=logging.DEBUG) | |
| unittest.main() |