"""Default runs learner functions. Add --reference to check the supplied engine.""" import copy import json import math from pathlib import Path import subprocess import sys import tempfile import unittest import tokenizer as engine if "--reference" in sys.argv: import tokenizer as solution sys.argv.remove("--reference") else: import starter as solution CORPUS = engine.default_corpus() def fit(docs, budget=16, checksum=""): return engine.train(docs, budget, checksum, solution.pair_counts, solution.merge_pair) def doc(text, doc_id="t1", split="train"): return {"id": doc_id, "split": split, "text": text} class Components(unittest.TestCase): def test_pair_counts_by_hand(self): self.assertEqual(dict(solution.pair_counts([[97, 98, 97, 98]])), {(97, 98): 2, (98, 97): 1}) def test_overlaps_count_but_replacement_does_not_overlap(self): self.assertEqual(solution.pair_counts([[97] * 4])[(97, 97)], 3) self.assertEqual(solution.merge_pair([97] * 4, (97, 97), 258), [258, 258]) self.assertEqual(solution.merge_pair([97] * 3, (97, 97), 258), [258, 97]) def test_documents_are_not_joined(self): self.assertEqual(dict(solution.pair_counts([[97], [98], [97], [98]])), {}) self.assertEqual(fit([doc("a", "1"), doc("b", "2"), doc("a", "3"), doc("b", "4")]).merges, ()) def test_no_input_mutation_or_global_replacement(self): ids = [1, 2, 1, 3, 1, 2] self.assertEqual(solution.merge_pair(ids, (1, 2), 258), [258, 1, 3, 258]) self.assertEqual(ids, [1, 2, 1, 3, 1, 2]) def test_empty_single_and_no_match(self): self.assertEqual(dict(solution.pair_counts([[], [7]])), {}) self.assertEqual(solution.merge_pair([], (1, 2), 258), []) self.assertEqual(solution.merge_pair([7], (1, 2), 258), [7]) def test_rank_order_builds_larger_tokens(self): ids = [97, 98, 97, 98] self.assertEqual(solution.apply_merges(ids, ((97, 98), (258, 258))), [259]) self.assertEqual(ids, [97, 98, 97, 98]) self.assertEqual(solution.apply_merges(ids, ()), ids) class Representation(unittest.TestCase): def setUp(self): self.model = fit([doc("abababab red red")]) def test_all_256_bytes_round_trip_without_unk(self): raw = bytes(range(256)) ids = self.model.encode_bytes(raw, solution.apply_merges) self.assertEqual(self.model.decode_bytes(ids), raw) self.assertTrue(set(ids).isdisjoint((engine.BOS, engine.EOS))) def test_unseen_scripts_emoji_code_and_whitespace_round_trip(self): for text in ("", "नमस्ते", "東京", "👩🏽‍💻", "\x00", "x = 1\n\t", "Zoë", "e\u0301"): with self.subTest(text=text): self.assertEqual(self.model.decode(self.model.encode(text, solution.apply_merges)), text) def test_no_silent_unicode_normalization(self): self.assertNotEqual(self.model.encode("é", solution.apply_merges), self.model.encode("e\u0301", solution.apply_merges)) def test_decode_bytes_before_decoding_text(self): self.assertEqual(self.model.decode([0xc3, 0xa9]), "é") self.assertEqual(self.model.decode_bytes([0xc3]), b"\xc3") with self.assertRaises(UnicodeDecodeError): self.model.decode([0xc3]) def test_reserved_spelling_is_literal_not_control(self): for text in ("", "", "<|endoftext|>"): ids = self.model.encode(text, solution.apply_merges) self.assertTrue(set(ids).isdisjoint((engine.BOS, engine.EOS))) self.assertEqual(self.model.decode(ids), text) def test_invalid_ids_types_and_encoding_are_rejected(self): for ids in ([engine.BOS], [engine.EOS], [-1], [999999], [True], ["97"]): with self.subTest(ids=ids), self.assertRaises(ValueError): self.model.decode_bytes(ids) for value in (None, 123, b"bytes"): with self.subTest(value=value), self.assertRaises(ValueError): self.model.encode(value) with self.assertRaises(UnicodeEncodeError): self.model.encode("\ud800") with self.assertRaises(ValueError): self.model.encode("a" * (engine.MAX_BYTES + 1)) class TrainingAndArtifacts(unittest.TestCase): def test_ties_and_training_order_are_deterministic(self): docs = [doc("baba", "1"), doc("abab", "2")] model = fit(docs, 1) self.assertEqual(model.merges, ((97, 98),)) self.assertEqual(model.merges, fit(list(reversed(docs)), 1).merges) self.assertEqual(model.digest(), fit(docs, 1).digest()) def test_fit_guards_budget_split_and_minimum_frequency(self): self.assertEqual(fit([doc("ab")]).merges, ()) self.assertEqual(fit([doc("abab")], 0).merges, ()) for budget in (-1, 129, True, 1.5): with self.subTest(budget=budget), self.assertRaises(ValueError): fit([doc("abab")], budget) for docs in ([], [doc("", "1")], [doc("abab", split="test")]): with self.subTest(docs=docs), self.assertRaises(ValueError): fit(docs) def test_save_load_round_trip_and_refuses_overwrite(self): model = fit([doc("abababab")]) with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "tokenizer.json" model.save(path) restored = engine.Tokenizer.load(path) self.assertEqual(restored.canonical_json(), model.canonical_json()) self.assertEqual(restored.encode("abab Ω"), model.encode("abab Ω")) with self.assertRaises(FileExistsError): model.save(path) def test_bad_artifact_version_policy_and_references_fail(self): payload = engine.Tokenizer().payload() bad = [{**payload, "format": "v999"}, {**payload, "policy": "lowercase"}, {**payload, "extra": True}, {**payload, "merges": [[258, 97]]}, {**payload, "merges": [[256, 97]]}, {**payload, "merges": [[97, 98], [97, 98]]}, {**payload, "merges": [[True, 98]]}, {**payload, "merges": "wrong"}, {**payload, "training_ids": ["x", "x"]}] for item in bad: with self.subTest(item=item), self.assertRaises(ValueError): engine.Tokenizer.from_payload(item) def test_expanding_artifact_is_bounded(self): pairs = [[97, 97]] + [[258 + i, 258 + i] for i in range(25)] with self.assertRaises(ValueError): engine.Tokenizer.from_payload({**engine.Tokenizer().payload(), "merges": pairs}) with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "large.json" path.write_bytes(b" " * 1048577) with self.assertRaises(ValueError): engine.Tokenizer.load(path) class IntegratedExperiment(unittest.TestCase): def test_real_learner_functions_drive_the_report(self): report = engine.run_experiment(CORPUS, solution.pair_counts, solution.merge_pair, solution.apply_merges) model = engine.Tokenizer.from_payload(report["tokenizer"]) self.assertEqual(report["tokenizer_sha256"], model.digest()) self.assertEqual(len(model.merges), 16) self.assertTrue(all(x["round_trip"] for x in report["representation"].values())) self.assertTrue(all(x["round_trip"] for x in report["transfer_probes"])) def test_changed_held_out_text_never_changes_fit(self): corpus, _, _ = engine.load_corpus(CORPUS) before = engine.run_experiment(CORPUS) changed = copy.deepcopy(corpus) changed["documents"][-1]["text"] = "Ω 👩🏽‍💻 entirely new test" changed["documents"][-3]["text"] = "新しい validation" with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "changed.json" path.write_text(json.dumps(changed), encoding="utf-8") after = engine.run_experiment(path) self.assertEqual(before["tokenizer"]["merges"], after["tokenizer"]["merges"]) self.assertEqual(before["tokenizer"]["training_ids"], after["tokenizer"]["training_ids"]) # Full corpus provenance hash changes, even though learned merges do not. self.assertNotEqual(before["corpus_sha256"], after["corpus_sha256"]) self.assertEqual(before["baselines"]["bpe"]["metrics"]["train"], after["baselines"]["bpe"]["metrics"]["train"]) def test_units_and_baseline_denominators_are_explicit(self): report = engine.run_experiment(CORPUS) _, splits, _ = engine.load_corpus(CORPUS) for kind, result in report["baselines"].items(): uniform = result["metrics"]["validation"]["uniform"] self.assertAlmostEqual(uniform["nll_nats_per_target"], math.log(result["output_vocabulary_size_including_eos"])) utf8_bytes = sum(len(d["text"].encode("utf-8")) for d in splits["validation"]) for metric in result["metrics"]["validation"].values(): self.assertAlmostEqual(metric["canonical_path_nats_per_utf8_byte"], metric["nll_nats_per_target"] * metric["targets_including_eos"] / utf8_bytes) if kind == "bpe": self.assertEqual(uniform["targets_including_eos"], report["representation"]["validation"]["content_tokens"] + len(splits["validation"])) self.assertFalse(report["baseline_contract"]["cross_tokenizer_token_perplexity_comparable"]) def test_corpus_rejects_duplicates_omissions_and_bad_types(self): corpus, _, _ = engine.load_corpus(CORPUS) variants = [] for field, value in (("text", corpus["documents"][0]["text"]), ("id", corpus["documents"][0]["id"]), ("split", [])): changed = copy.deepcopy(corpus) changed["documents"][-1][field] = value variants.append(changed) variants.append({**corpus, "documents": [d for d in corpus["documents"] if d["split"] != "test"]}) variants.append({**corpus, "documents": [None]}) for changed in variants: with self.subTest(), tempfile.TemporaryDirectory() as directory: path = Path(directory) / "bad.json" path.write_text(json.dumps(changed), encoding="utf-8") with self.assertRaises(ValueError): engine.load_corpus(path) def test_cli_report_is_repeatable_and_cannot_overwrite(self): script = Path(engine.__file__) with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "report.json" args = [sys.executable, "-B", str(script), "--corpus", str(CORPUS), "--output", str(path)] result = subprocess.run(args, capture_output=True, text=True, timeout=10) self.assertEqual(result.returncode, 0, result.stderr) original = path.read_bytes() self.assertEqual(json.loads(original), engine.run_experiment(CORPUS)) self.assertNotEqual(subprocess.run(args, capture_output=True, timeout=10).returncode, 0) self.assertEqual(path.read_bytes(), original) if __name__ == "__main__": unittest.main()