Something went wrong. Try again.
a tiny pet with an even tinier brain, living in your browser
Something went wrong. Try again.
27 kB · 464 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465"""Stdlib acceptance regressions and optional isolated-environment Torch rules."""import copyimport importlib.utilimport ioimport mathfrom pathlib import Pathimport sysimport unittestfrom contextlib import redirect_stderrfrom unittest.mock import patch
ROOT = Path(__file__).resolve().parent.parentspec = importlib.util.spec_from_file_location("experiments", ROOT / "train/experiments.py")E = importlib.util.module_from_spec(spec)spec.loader.exec_module(E)
TRAINING_AVAILABLE = all(importlib.util.find_spec(name) is not None for name in ("torch", "transformers", "tokenizers", "huggingface_hub"))if TRAINING_AVAILABLE: import torch from tokenizers import Regex, Tokenizer, decoders, models, pre_tokenizers from transformers import GPTNeoXConfig, GPTNeoXForCausalLM, PreTrainedTokenizerFast sys.path.insert(0, str(ROOT / "train")) try: train_spec = importlib.util.spec_from_file_location("pet_train", ROOT / "train/train.py") T = importlib.util.module_from_spec(train_spec) train_spec.loader.exec_module(T) prepare_spec = importlib.util.spec_from_file_location("prepare_model", ROOT / "train/prepare_model.py") P = importlib.util.module_from_spec(prepare_spec) prepare_spec.loader.exec_module(P) eval_spec = importlib.util.spec_from_file_location("eval_pet", ROOT / "train/eval_pet.py") V = importlib.util.module_from_spec(eval_spec) eval_spec.loader.exec_module(V) finally: sys.path.pop(0)
class AcceptanceRules(unittest.TestCase): def setUp(self): self.policy = {"minimum": ["valid_lines"], "maximum": ["healthy_medicine_yes"], "margin": 0} self.reference = {"status": "complete", "dataset_sha256": "corpus-a", "scenario_sha256": "fixed-scripted-cases", "experiment_contract": {"seed": 0, "samples": 40}, "rates": {"valid_lines": 0.9, "healthy_medicine_yes": 0.1}, "rate_counts": {"valid_lines": {"numerator": 90, "denominator": 100}, "healthy_medicine_yes": {"numerator": 10, "denominator": 100}}}
def candidate(self, metric, numerator, denominator=100): report = copy.deepcopy(self.reference) report["rates"][metric] = numerator / denominator if denominator else None report["rate_counts"][metric] = {"numerator": numerator, "denominator": denominator} return report
def gate(self, candidate): return E.quality_gate(self.reference, candidate, self.policy)["passed"]
def test_accepts_equal_or_improved_behavior(self): self.assertTrue(self.gate(self.reference)) self.assertTrue(self.gate(self.candidate("valid_lines", 91))) self.assertTrue(self.gate(self.candidate("healthy_medicine_yes", 9)))
def test_rejects_one_success_lost_or_one_unsafe_offer_added(self): self.assertFalse(self.gate(self.candidate("valid_lines", 89))) self.assertFalse(self.gate(self.candidate("healthy_medicine_yes", 11)))
def test_rejects_missing_or_zero_denominator(self): report = copy.deepcopy(self.reference) del report["rate_counts"]["valid_lines"] self.assertFalse(self.gate(report)) self.assertFalse(self.gate(self.candidate("valid_lines", 0, 0)))
def test_rejects_smaller_cherry_picked_suite(self): self.assertFalse(self.gate(self.candidate("valid_lines", 9, 10)))
def test_rejects_mismatched_corpus_or_scenario_contract(self): for key, value in (("dataset_sha256", "other-data"), ("scenario_sha256", "changed-cases"), ("experiment_contract", {"seed": 1, "samples": 40}), ("config", {"device": "cuda"})): report = copy.deepcopy(self.reference) report[key] = value self.assertFalse(self.gate(report))
def test_rejects_incomplete_reference_or_candidate(self): report = copy.deepcopy(self.reference) report["status"] = "failed" self.assertFalse(self.gate(report)) self.reference["status"] = "failed" self.assertFalse(self.gate(report))
def test_rejects_nan_or_inconsistent_rates(self): for value in (float("nan"), 0.95, True): report = copy.deepcopy(self.reference) report["rates"]["valid_lines"] = value self.assertFalse(self.gate(report))
def test_pareto_preserves_size_speed_tradeoff(self): small = {"id": "small", "parameters": 9, "graph_bytes": 5, "sampled_turn_ms": 30} fast = {"id": "fast", "parameters": 13, "graph_bytes": 7, "sampled_turn_ms": 10} dominated = {"id": "worse", "parameters": 14, "graph_bytes": 8, "sampled_turn_ms": 40} self.assertEqual(E.pareto_frontier([small, fast, dominated]), [small, fast])
def test_runtime_gate_requires_complete_real_turn_timings(self): report = {"status": "complete", "graph_sha256": "measured", "graph_bytes": 123, "browser_version": "browser", "device": "wasm", "cases": [{"status": "complete", "cross_origin_isolated": True, "config": {"threads": 1, "seed": 0}, "result": { "config": {"prompts": "both", "mode": "both", "repeats": 3, "seed": 0}, "results": [{"scenario": scenario, "mode": mode} for scenario in ("short_idle", "long_react") for mode in ("greedy", "sampled")], "summary": {"sampled_turn_ms": 200}}}]} self.assertEqual(E.runtime_score(report)["sampled_turn_ms"], 200) report["cases"][0]["result"]["summary"]["sampled_turn_ms"] = float("inf") with self.assertRaises(ValueError): E.runtime_score(report) report["cases"][0]["result"]["summary"]["sampled_turn_ms"] = 200 report["cases"][0]["status"] = "failed" with self.assertRaises(ValueError): E.runtime_score(report) good = copy.deepcopy(report["cases"][0]) good.update(status="complete", config={"threads": 4, "qdq": 0, "seed": 0}) good["result"]["summary"]["sampled_turn_ms"] = 100 report["cases"].append(good) self.assertEqual(E.runtime_score(report, {"qdq": 0})["sampled_turn_ms"], 100) with self.assertRaises(ValueError): E.runtime_score(report, {"qdq": 1}) good["result"]["config"]["prompts"] = "short" with self.assertRaises(ValueError): E.runtime_score(report, {"qdq": 0})
class BenchmarkRules(unittest.TestCase): def setUp(self): names = {name for metrics in E.CAPABILITY_METRICS.values() for name in metrics} self.report = {"status": "complete", "rates": {name: 0.5 for name in names}, "rate_counts": {name: {"numerator": 5, "denominator": 10} for name in names}} self.contract = {"version": 2, "seed": 0, "samples": 40} self.report.update(backend="onnx", dataset_sha256="dataset", experiment_contract=self.contract, scenario_sha256="scenarios", scenario_suite_sha256={"core": "core-cases"}, scenario_counts={"core": 40}, config={"seed": 0, "device": "cpu", "sampling": {"temperature": 0.85, "eos_token_ids": [0]}, "generation": {"react": 64}}, onnx_artifact={"graph_sha256": "graph", "dtype": "q4", "disable_quant_qdq": True, "disable_prepacking": True})
def test_balanced_capabilities_use_counts_without_rewarding_safety(self): score = E.capability_index(self.report) self.assertEqual(score["quality"], 50) self.report["rates"]["quiet_idle"] = 1 self.report["rate_counts"]["quiet_idle"]["numerator"] = 10 self.report["rates"]["unsafe_diaries"] = 0 self.assertAlmostEqual(E.capability_index(self.report)["quality"], 100 * 3.5 / 6) self.assertEqual(score["categories"]["memory"]["metrics"]["memory_saved"]["denominator"], 10)
def test_missing_or_zero_denominator_never_becomes_perfect_or_zero_score(self): for name in ("memory_saved", "diary_format_valid"): report = copy.deepcopy(self.report) del report["rate_counts"][name] with self.assertRaises(ValueError): E.capability_index(report) report["rate_counts"][name] = {"numerator": 0, "denominator": 0} report["rates"][name] = 0 with self.assertRaises(ValueError): E.capability_index(report)
def test_mismatched_actual_graph_scenario_data_seed_or_contract_rejected(self): E.validate_benchmark_behavior(self.report, "graph", "q4", "dataset", self.contract, self.report) mutations = (("onnx_artifact", dict(self.report["onnx_artifact"], graph_sha256="other")), ("scenario_sha256", "other"), ("dataset_sha256", "other"), ("experiment_contract", dict(self.contract, seed=1)), ("config", dict(self.report["config"], seed=1))) for key, value in mutations: report = copy.deepcopy(self.report) report[key] = value with self.assertRaises(ValueError, msg=key): E.validate_benchmark_behavior(report, "graph", "q4", "dataset", self.contract, self.report)
def row(self, name, quality=50, speed=100, size=10, status="complete"): return {"id": name, "dtype": "q4", "quality": quality, "browser_fixed_decode_tps": speed, "graph_bytes": size, "status": status, "failures": []}
def test_failed_and_missing_rows_are_unranked(self): failed = self.row("failed", quality=100, speed=1000, status="failed") missing = self.row("missing", quality=None) incomplete = self.row("incomplete", speed=None) error = self.row("error", quality=100) error["failures"] = [{"phase": "native", "error": "failed"}] good = self.row("good") E.rank_benchmark([failed, missing, incomplete, error, good]) for row in (failed, missing, incomplete, error): self.assertIsNone(row["rank"]) self.assertIsNone(row["speed_rank"]) self.assertEqual(good["rank"], 1)
def test_quality_then_fixed_speed_then_size_tie_precedence(self): quality = self.row("quality", quality=51, speed=1, size=100) speed = self.row("speed", speed=101, size=100) small = self.row("small", size=9) large = self.row("large", size=10) rows = [large, small, speed, quality] E.rank_benchmark(rows) self.assertEqual([r["id"] for r in sorted(rows, key=lambda r: r["rank"])], ["quality", "speed", "small", "large"]) self.assertEqual([r["id"] for r in sorted(rows, key=lambda r: r["speed_rank"])], ["speed", "small", "large", "quality"])
@unittest.skipUnless(TRAINING_AVAILABLE, "requires isolated torch/transformers/tokenizers training environment")class MaskedDistillationRules(unittest.TestCase): def tokenizers(self, swap=False, tokenizer_class=None): vocab = {"<eos>": 0, "<unk>": 1, "a": 3 if swap else 2, "b": 2 if swap else 3} backend = Tokenizer(models.WordLevel(vocab=vocab, unk_token="<unk>")) backend.pre_tokenizer = pre_tokenizers.Split(Regex("."), behavior="isolated") backend.decoder = decoders.Fuse() cls = tokenizer_class or PreTrainedTokenizerFast return cls(tokenizer_object=backend, eos_token="<eos>", pad_token="<eos>", unk_token="<unk>")
def logits(self): labels = torch.tensor([[-100, -100, 1, -100, 2]]) student_prob = torch.tensor([[0.2, 0.3, 0.5], [0.2, 0.3, 0.5]]) teacher_prob = torch.tensor([[0.5, 0.3, 0.2], [0.1, 0.2, 0.7]]) student, teacher = torch.zeros(1, 5, 3), torch.zeros(1, 5, 3) student[0, [1, 3]] = 2 * student_prob.log() teacher[0, [1, 3]] = 2 * teacher_prob.log() return student, teacher, labels, student_prob, teacher_prob
def test_masked_kl_matches_expected_probabilities_and_temperature_formula(self): student, teacher, labels, student_prob, teacher_prob = self.logits() mask = labels[:, 1:] != -100 student_targets = T.supervised_logits(student, mask) teacher_targets = T.supervised_logits(teacher, mask) self.assertEqual(student_targets.shape, (2, 3)) kl = T.distillation_kl(student_targets, teacher_targets, 2) expected = (teacher_prob * (teacher_prob.log() - student_prob.log())).sum(-1).mean() torch.testing.assert_close(kl, expected, rtol=1e-6, atol=1e-7) ce = torch.nn.functional.cross_entropy(student_targets.float(), labels[:, 1:][mask]) weight = 0.35 total = (1 - weight) * ce + weight * 2 ** 2 * kl expected_ce = -(2 * student_prob.log()).log_softmax(-1)[[0, 1], [1, 2]].mean() torch.testing.assert_close(total, (1 - weight) * expected_ce + weight * 4 * expected, rtol=1e-6, atol=1e-7)
def test_masked_positions_and_final_logits_cannot_change_kl_or_receive_gradients(self): student, teacher, labels, _, _ = self.logits() mask = labels[:, 1:] != -100 expected = T.distillation_kl(T.supervised_logits(student, mask), T.supervised_logits(teacher, mask), 2) # Headers/owner/event/results/padding are represented only by -100 targets. student[0, [0, 2, 4]] = torch.tensor([1000., -1000., 1000.]) teacher[0, [0, 2, 4]] = torch.tensor([-1000., 1000., -1000.]) student.requires_grad_(True) teacher.requires_grad_(True) with torch.enable_grad(): actual = T.distillation_kl(T.supervised_logits(student, mask), T.supervised_logits(teacher, mask), 2) torch.testing.assert_close(actual, expected, rtol=0, atol=0) actual.backward() self.assertTrue(torch.is_grad_enabled()) self.assertIsNone(teacher.grad) self.assertEqual(torch.count_nonzero(student.grad[0, [0, 2, 4]]).item(), 0) self.assertGreater(torch.count_nonzero(student.grad[0, [1, 3]]).item(), 0)
def test_bfloat16_targets_use_fp32_numerical_kl(self): student, teacher, labels, _, _ = self.logits() mask = labels[:, 1:] != -100 student_targets = T.supervised_logits(student.to(torch.bfloat16), mask).requires_grad_(True) teacher_targets = T.supervised_logits(teacher.to(torch.bfloat16), mask).requires_grad_(True) kl = T.distillation_kl(student_targets, teacher_targets, 2) self.assertEqual(kl.dtype, torch.float32) self.assertTrue(math.isfinite(kl.item())) kl.backward() self.assertIsNotNone(student_targets.grad) self.assertIsNone(teacher_targets.grad)
def test_rejects_same_size_but_different_token_to_id_mapping(self): with self.assertRaisesRegex(ValueError, "teacher/student tokenizer mismatch"): T.validate_teacher_tokenizers(self.tokenizers(), self.tokenizers(swap=True))
def test_identical_tokenizers_encode_all_full_text_before_truncation(self): student, teacher = self.tokenizers(), self.tokenizers() match = T.validate_teacher_tokenizers(student, teacher) self.assertEqual(match["student"]["serialized_sha256"], match["teacher"]["serialized_sha256"]) row = {"kind": "chat", "text": "ababa", "spans": [[1, 5]]} encoded, stats = T.encode_corpus(student, [row], 2, "pet", "training", teacher) self.assertEqual(encoded, [([2, 3], [-100, 3])]) self.assertEqual(stats["supervised_tokens_full"], 4) self.assertEqual(stats["supervised_tokens_retained"], 1)
def test_rejects_actual_encoding_mismatch_even_beyond_cap_and_on_heldout(self): class DifferentEncoding(PreTrainedTokenizerFast): def __call__(self, *args, **kwargs): encoded = super().__call__(*args, **kwargs) encoded["input_ids"][-1] = self.get_vocab()["b"] return encoded student, teacher = self.tokenizers(), self.tokenizers(tokenizer_class=DifferentEncoding) T.validate_teacher_tokenizers(student, teacher) row = {"kind": "chat", "text": "ababa", "spans": [[1, 5]]} for split in ("training", "held-out"): with self.assertRaisesRegex(ValueError, split + " row 1: teacher/student encoding mismatch"): T.encode_corpus(student, [row], 2, "pet", split, teacher)
def test_rejects_invalid_distillation_weights_temperatures_and_stage(self): required = ["train.py", "--data", "unused.jsonl", "--base", "unused", "--out", "unused-out"] invalid = ((["--distill-weight", "nan"], "distill-weight must be finite"), (["--distill-weight", "-0.1"], "distill-weight must be finite"), (["--distill-weight", "1.1"], "distill-weight must be finite"), (["--distill-temperature", "inf"], "distill-temperature must be finite"), (["--distill-temperature", "0"], "distill-temperature must be finite"), (["--distill-temperature", "-1"], "distill-temperature must be finite"), (["--distill-weight", "0.5"], "positive distill-weight requires --teacher")) for flags, message in invalid: with self.subTest(flags=flags), patch.object(sys, "argv", required + flags), redirect_stderr(io.StringIO()) as err: with self.assertRaises(SystemExit) as exit: T.main() self.assertEqual(exit.exception.code, 2) self.assertIn(message, err.getvalue()) pretrain = ["train.py", "--pretrain-data", "unused.jsonl", "--base", "unused", "--out", "unused-out", "--teacher", str(ROOT), "--distill-weight", "0.5"] with patch.object(sys, "argv", pretrain), redirect_stderr(io.StringIO()) as err: with self.assertRaises(SystemExit) as exit: T.main() self.assertEqual(exit.exception.code, 2) self.assertIn("only supported with --data pet rows", err.getvalue())
def test_rejects_empty_or_mismatched_target_logits(self): for student, teacher in ((torch.zeros(0, 3), torch.zeros(0, 3)), (torch.zeros(2, 3), torch.zeros(2, 4))): with self.assertRaisesRegex(ValueError, "equal nonempty"): T.distillation_kl(student, teacher, 1) with self.assertRaisesRegex(ValueError, "shifted boolean target mask"): T.supervised_logits(torch.zeros(1, 4, 3), torch.zeros(1, 4, dtype=torch.bool))
@unittest.skipUnless(TRAINING_AVAILABLE, "requires isolated torch/transformers/tokenizers training environment")class EvaluatorRules(unittest.TestCase): def test_lookup_counts_survive_malformed_outer_tag(self): from types import SimpleNamespace evaluator = V.Evaluator(None, SimpleNamespace(eos_token_id=0), torch.device('cpu'), 0, config=SimpleNamespace(max_position_embeddings=1024)) pieces = iter([('[remember:bad [clock]', 'lookup'), (' [clock]', 'lookup'), (' hi', 'line')]) def complete(*args): text, reason = next(pieces) return {'success': True, 'text': text, 'reason': reason} evaluator.complete = complete state = V.G.rand_state(V.G.random.Random(737)) state['asleep'] = False result = evaluator.pet_line(state, [], 'react', evaluator.generator('case')) self.assertEqual(result['lookups_requested'], 2) self.assertEqual(result['lookups_resolved'], 2) self.assertEqual(result['lookups_unresolved'], 0) self.assertFalse(V.valid(result['text']))
def test_small_context_does_not_hide_production_prompt_overflow(self): class CharacterTokenizer: def encode(self, text, **kwargs): return list(text) with self.assertRaises(V.ContextOverflow): V.fit_prompt(CharacterTokenizer(), 'H', ['x' * 30, 'z' * 10], 'P:', 10, 1024, 40) _, details = V.fit_prompt(CharacterTokenizer(), 'H', ['x' * 30, 'z' * 10], 'P:', 10, 40, 40) self.assertEqual(details['kept_lines'], 1)
@unittest.skipUnless(TRAINING_AVAILABLE, "requires isolated torch/transformers/tokenizers training environment")class CompressionRules(unittest.TestCase): def model(self, vocab_size=32): config = GPTNeoXConfig(vocab_size=vocab_size, hidden_size=16, intermediate_size=32, num_hidden_layers=6, num_attention_heads=4, max_position_embeddings=32, rotary_pct=0.5, hidden_dropout=0, attention_dropout=0, tie_word_embeddings=False) config._attn_implementation = "eager" return GPTNeoXForCausalLM(config).eval()
def test_pruning_retains_exact_blocks_and_three_layer_greedy_cache_matches_fresh(self): torch.manual_seed(7) model = self.model() indices = [0, 2, 5] original_blocks = list(model.gpt_neox.layers) states = [{name: value.clone() for name, value in original_blocks[index].state_dict().items()} for index in indices] before = sum(p.numel() for p in model.parameters()) removed = sum(p.numel() for index, block in enumerate(original_blocks) if index not in indices for p in block.parameters()) inp = model.get_input_embeddings().weight.detach().clone() out = model.get_output_embeddings().weight.detach().clone() receipt = P.select_layers(model, indices) self.assertEqual(model.config.num_hidden_layers, 3) self.assertEqual(receipt["source_indices"], indices) self.assertEqual(receipt["cache_indices"], [0, 1, 2]) self.assertEqual(sum(p.numel() for p in model.parameters()), before - removed) for position, block in enumerate(model.gpt_neox.layers): self.assertIs(block, original_blocks[indices[position]]) self.assertEqual(block.attention.layer_idx, position) for name, value in block.state_dict().items(): torch.testing.assert_close(value, states[position][name], rtol=0, atol=0) torch.testing.assert_close(model.get_input_embeddings().weight, inp, rtol=0, atol=0) torch.testing.assert_close(model.get_output_embeddings().weight, out, rtol=0, atol=0) with torch.no_grad(): ids = torch.tensor([[1, 2, 3]]) cached = model(input_ids=ids, use_cache=True) for _ in range(3): self.assertEqual(len(cached.past_key_values), 3) self.assertEqual(cached.past_key_values.get_seq_length(), ids.shape[1]) next_id = cached.logits[:, -1].argmax(-1, keepdim=True) ids = torch.cat([ids, next_id], dim=1) cached = model(input_ids=next_id, past_key_values=cached.past_key_values, use_cache=True) fresh = model(input_ids=ids, use_cache=False) torch.testing.assert_close(cached.logits[:, -1], fresh.logits[:, -1], rtol=1e-5, atol=1e-6) self.assertTrue(torch.equal(cached.logits[:, -1].argmax(-1), fresh.logits[:, -1].argmax(-1)))
def test_domain_selection_decomposes_source_added_whitespace_into_genuine_bpe(self): alphabet = sorted(pre_tokenizers.ByteLevel.alphabet()) eos, unused = "<eos>", " " # The source tokenizer can shortcut spaces to an added row without BPE # ancestors. Domain frequencies must use the actual adapted backend. pieces = alphabet + [unused, "ab", "bc", "xy", "xyz", eos] vocab = {piece: index for index, piece in enumerate(pieces)} merges = [("a", "b"), ("b", "c"), ("x", "y"), ("xy", "z")] tok = P.wrap_tokenizer(P.bytelevel_backend(vocab, merges), eos, 32) tok.add_tokens([unused]) self.assertIn(vocab[unused], tok.encode('xyz xyz', add_special_tokens=False)) selected, retained, source_ids = P.domain_vocabulary(vocab, merges, eos, len(vocab) - 1, [{"text": "xyz xyz xyz"}]) self.assertEqual(set(selected), set(vocab) - {unused}) self.assertEqual(retained, merges) self.assertNotIn(vocab[unused], source_ids) self.assertEqual(source_ids[selected[eos]], vocab[eos]) with self.assertRaisesRegex(ValueError, "invalid domain vocabulary"): P.domain_vocabulary(vocab, merges, "<missing-eos>", len(vocab) - 1, [{"text": "xyz"}])
def test_domain_selection_is_nonprefix_merge_closed_and_preserves_real_eos_and_both_tables(self): alphabet = sorted(pre_tokenizers.ByteLevel.alphabet()) eos = "<eos>" pieces = alphabet + ["ab", "bc", "xy", "xyz", eos] vocab = {piece: index for index, piece in enumerate(pieces)} merges = [("a", "b"), ("b", "c"), ("x", "y"), ("xy", "z")] selected, retained, source_ids = P.domain_vocabulary(vocab, merges, eos, len(alphabet) + 3, [{"text": "xyz xyz xyz xyz"}]) self.assertEqual(len(selected), len(alphabet) + 3) self.assertIn("xyz", selected) self.assertIn("xy", selected) self.assertNotIn("ab", selected) self.assertNotIn("bc", selected) self.assertEqual(retained, [("x", "y"), ("xy", "z")]) for a, b in retained: self.assertTrue({a, b, a + b} <= selected.keys()) self.assertTrue(P.UTF8_BYTE_SYMBOLS <= selected.keys()) self.assertEqual(source_ids[selected[eos]], vocab[eos]) self.assertNotEqual(selected[eos], vocab[eos]) self.assertEqual(source_ids, [vocab[piece] for piece, index in sorted(selected.items(), key=lambda item: item[1])]) reduced_tok = P.wrap_tokenizer(P.bytelevel_backend(selected, retained), eos, 32) self.assertEqual(reduced_tok(eos, add_special_tokens=False)["input_ids"], [selected[eos]]) self.assertEqual(reduced_tok.decode(reduced_tok("xyz xyz", add_special_tokens=False)["input_ids"], clean_up_tokenization_spaces=False), "xyz xyz") model = self.model(vocab_size=len(vocab)) before = sum(p.numel() for p in model.parameters()) input_rows = model.get_input_embeddings().weight.detach()[source_ids].clone() output_rows = model.get_output_embeddings().weight.detach()[source_ids].clone() body_hash = P.backbone_hash(model) receipt = P.reduce_pythia(model, source_ids) torch.testing.assert_close(model.get_input_embeddings().weight, input_rows, rtol=0, atol=0) torch.testing.assert_close(model.get_output_embeddings().weight, output_rows, rtol=0, atol=0) self.assertIsNot(model.get_input_embeddings().weight, model.get_output_embeddings().weight) self.assertEqual(receipt["new_id_to_source_id"][selected[eos]], vocab[eos]) self.assertEqual(P.backbone_hash(model), body_hash) self.assertEqual(sum(p.numel() for p in model.parameters()), before - 2 * (len(vocab) - len(selected)) * 16)
if __name__ == "__main__": unittest.main()