test_phase5.py

"""
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)