""" Phase 5 regression test suite. NO model load. Tests: format invariants, token_f1, generate (mock), dedup, retrieval gating. Runtime target: <10s. """ import sys, shutil import torch # ── Mock tokenizer ──────────────────────────────────────────────────────── class _Enc: def __init__(self, ids): self.ids = ids class MockTok: """char → ord(). Deterministic, injective, no special-token injection.""" def encode(self, text, add_special_tokens=True): return _Enc([ord(c) for c in text]) # ── Imports (NO model) ──────────────────────────────────────────────────── from eval_benchmark import ( clean_ids, build_prompt_ids, token_f1, generate, BOS_ID, USER_ID, ASSISTANT_ID, EOS_ID, MAX_GEN, ) from online_learner import build_update_seq, LORA_ALPHA, LORA_RANK from task_memory import TaskMemory from lora import LoRAAdapter tok = MockTok() passed = 0 failed = 0 def check(name, cond, msg=''): global passed, failed if cond: print(f' PASS {name}') passed += 1 else: print(f' FAIL {name}: {msg}') failed += 1 # ════════════════════════════════════════════════════════════════════════════ # 1. clean_ids — no special tokens ever appear # ════════════════════════════════════════════════════════════════════════════ ids = clean_ids(tok, 'hello') check('clean_ids_no_BOS', BOS_ID not in ids, f'BOS in {ids}') check('clean_ids_no_EOS', EOS_ID not in ids, f'EOS in {ids}') check('clean_ids_deterministic', ids == clean_ids(tok, 'hello')) check('clean_ids_empty', clean_ids(tok, '') == []) # ════════════════════════════════════════════════════════════════════════════ # 2. build_prompt_ids — exact canonical format # ════════════════════════════════════════════════════════════════════════════ q = 'Who?' p = build_prompt_ids(tok, q) qids = clean_ids(tok, q) check('prompt_starts_BOS_USER', p[:2] == [BOS_ID, USER_ID], f'prefix={p[:2]}') check('prompt_ends_ASSISTANT', p[-1] == ASSISTANT_ID, f'suffix={p[-1]}') check('prompt_exact', p == [BOS_ID, USER_ID] + qids + [ASSISTANT_ID]) check('prompt_no_extra_specials', p.count(BOS_ID) == 1 and p.count(EOS_ID) == 0) # ════════════════════════════════════════════════════════════════════════════ # 3. build_update_seq — label alignment invariant # ════════════════════════════════════════════════════════════════════════════ q2, a2 = 'Who invented telephone?', 'Bell' x_t, y_t = build_update_seq(tok, q2, a2) x = x_t[0].tolist() lbl = y_t[0].tolist() prompt2 = build_prompt_ids(tok, q2) answer_ids = clean_ids(tok, a2) # no EOS n_masked = len(prompt2) - 1 # positions where label=-100 check('update_seq_len_match', len(x) == len(lbl), f'x={len(x)} lbl={len(lbl)}') check('labels_prefix_masked', all(v == -100 for v in lbl[:n_masked]), f'not all -100 in first {n_masked} positions') check('labels_answer_body', lbl[n_masked:n_masked + len(answer_ids)] == answer_ids, f'lbl answer portion mismatch') check('labels_ends_EOS', lbl[n_masked + len(answer_ids)] == EOS_ID, f'label after answer = {lbl[n_masked + len(answer_ids)]}') # x = full[:-1]: positions [n_prompt:] == clean(answer) (EOS stripped by [:-1]) n_prompt = len(prompt2) check('x_answer_portion', x[n_prompt:n_prompt + len(answer_ids)] == answer_ids, f'x[{n_prompt}:{n_prompt+len(answer_ids)}] mismatch') # ════════════════════════════════════════════════════════════════════════════ # 4. token_f1 — metric correctness # ════════════════════════════════════════════════════════════════════════════ check('f1_perfect', abs(token_f1([1,2,3], [1,2,3]) - 1.0) < 1e-9) check('f1_zero', token_f1([1,2,3], [4,5,6]) == 0.0) check('f1_empty_pred', token_f1([], [1,2]) == 0.0) check('f1_empty_ref', token_f1([1,2], []) == 0.0) # F1([1,2,3,4],[1,2,5,6]): common=2, P=2/4=0.5, R=2/4=0.5 → F1=0.5 check('f1_half', abs(token_f1([1,2,3,4], [1,2,5,6]) - 0.5) < 1e-9, f'got {token_f1([1,2,3,4],[1,2,5,6])}') # ════════════════════════════════════════════════════════════════════════════ # 5. generate — EOS terminates output; output never contains EOS # ════════════════════════════════════════════════════════════════════════════ class _MockModel: """Outputs token 42 until seq len >= 5, then outputs EOS.""" def __call__(self, ids, adapter=None): T = ids.shape[1] V = max(EOS_ID + 1, 100) logits = torch.zeros(1, T, V) if T >= 5: logits[0, -1, EOS_ID] = 100.0 else: logits[0, -1, 42] = 100.0 return logits gen = generate(_MockModel(), tok, [BOS_ID, USER_ID, 100], max_new=MAX_GEN) # T starts=3: outputs 42,42 then EOS at T=5 → stops; gen=[42,42] check('generate_no_EOS_in_output', EOS_ID not in gen, f'EOS found: {gen}') check('generate_body_correct', all(t == 42 for t in gen), f'gen={gen}') check('generate_within_max', len(gen) <= MAX_GEN) # max_new=0 → empty gen_zero = generate(_MockModel(), tok, [BOS_ID], max_new=0) check('generate_max_new_zero', gen_zero == []) # ════════════════════════════════════════════════════════════════════════════ # 6. LORA_ALPHA / LORA_RANK cross-file invariant # ════════════════════════════════════════════════════════════════════════════ check('alpha_cross_file', LORA_ALPHA == TaskMemory.ALPHA, f'online={LORA_ALPHA} mem={TaskMemory.ALPHA}') check('rank_cross_file', LORA_RANK == TaskMemory.RANK, f'online={LORA_RANK} mem={TaskMemory.RANK}') # ════════════════════════════════════════════════════════════════════════════ # 7. TaskMemory dedup — same embedding → overwrite, memory bounded # ════════════════════════════════════════════════════════════════════════════ TEST_DIR = '/tmp/tm_regtest_phase5' shutil.rmtree(TEST_DIR, ignore_errors=True) mem = TaskMemory(TEST_DIR, top_k=3) emb_a = torch.randn(1, 512) emb_b = torch.randn(1, 512) ad1 = LoRAAdapter(8, 512, 4, 32.0) ad2 = LoRAAdapter(8, 512, 4, 32.0) ad3 = LoRAAdapter(8, 512, 4, 32.0) tid0 = mem.add(emb_a, ad1, {'post_score': 0.6}) tid1 = mem.add(emb_b, ad2, {'post_score': 0.7}) tid0b = mem.add(emb_a, ad3, {'post_score': 0.9}) # same emb_a → dedup check('dedup_task_id_reused', tid0b == tid0, f'expected {tid0}, got {tid0b}') check('dedup_memory_bounded', len(mem) == 2, f'len={len(mem)}') check('dedup_faiss_bounded', mem.index.ntotal == 2, f'ntotal={mem.index.ntotal}') check('dedup_meta_updated', mem.metadata[tid0]['post_score'] == 0.9, str(mem.metadata[tid0])) # Adapter file must contain ad3's weights mem.flush() ckpt = torch.load(mem._adapter_path(tid0), map_location='cpu', weights_only=True) sd3 = ad3.state_dict() check('dedup_adapter_overwritten', all(torch.equal(ckpt['state_dict'][k], sd3[k]) for k in sd3)) # ════════════════════════════════════════════════════════════════════════════ # 8. DIST_THRESHOLD — far embedding → no retrieval; close → retrieval # ════════════════════════════════════════════════════════════════════════════ # squared L2 from emb_a to far_emb >> DIST_THRESHOLD=5.0 far_emb = torch.zeros(1, 512) far_emb[0, 0] = 1000.0 adapters, _ = mem.retrieve(far_emb) check('dist_threshold_blocks_far', len(adapters) == 0, f'got {len(adapters)} adapters') # emb_a is an exact stored key → must retrieve merged = mem.retrieve_merged(emb_a) check('retrieve_merged_close', merged is not None) shutil.rmtree(TEST_DIR, ignore_errors=True) # ════════════════════════════════════════════════════════════════════════════ # 9. TaskMemory persistence — reload from disk restores state # ════════════════════════════════════════════════════════════════════════════ TEST_DIR2 = '/tmp/tm_persist_phase5' shutil.rmtree(TEST_DIR2, ignore_errors=True) mem2 = TaskMemory(TEST_DIR2, top_k=3) emb_c = torch.randn(1, 512) ad4 = LoRAAdapter(8, 512, 4, 32.0) mem2.add(emb_c, ad4, {'post_score': 0.8}) mem2.flush() del mem2 mem2b = TaskMemory(TEST_DIR2, top_k=3) check('persist_len', len(mem2b) == 1) check('persist_meta', mem2b.metadata[0]['post_score'] == 0.8) merged2 = mem2b.retrieve_merged(emb_c) check('persist_retrieve', merged2 is not None) shutil.rmtree(TEST_DIR2, ignore_errors=True) # ════════════════════════════════════════════════════════════════════════════ print(f'\n{passed}/{passed + failed} tests passed') if failed > 0: sys.exit(1)