| # 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() |
