#!/usr/bin/env python3 """Phase 7.4 — Zero-forgetting benchmark: KAIZEN episodic vs shared-adapter. CLAIM: per-task adapter isolation = zero catastrophic forgetting by construction. Learning task B never modifies adapter_A. Storage is fully isolated. This benchmark validates BOTH: (1) Storage isolation: after training domains 2..5, domain 1 adapters unchanged → domain 1 retention stays at its peak (KAIZEN BWT ≈ 0) (2) Shared adapter degrades: overwriting one adapter for all tasks causes forgetting of earlier domains (SHARED BWT < 0) Domain design: 5 groups × 10 tasks from the validated BENCHMARK_TASKS set. BENCHMARK_TASKS were already confirmed to give high inter-task embedding distance (Phase 5: F1→0.99 with 50 tasks, bounded memory, no cross-task contamination). This guarantees retrieval purity across domains. Test = SAME training questions (retention test, not generalization test). KAIZEN's episodic retrieval correctly returns exact stored adapters for exact question repeats (d≈0.0 → guaranteed below DIST_THRESHOLD=5.0). """ import time import torch from tokenizers import Tokenizer from huggingface_hub import hf_hub_download from lora import KaizenWithLoRA, LoRAAdapter from task_memory import TaskMemory from eval_benchmark import BENCHMARK_TASKS, build_prompt_ids, clean_ids, generate, token_f1, BLOCK_SIZE, MAX_GEN from online_learner import ( HF_TOKEN, HF_TOK_REPO, TOP_K, LORA_RANK, LORA_ALPHA, ATTEMPT_THRESHOLD, STORE_THRESHOLD, build_update_seq, online_update, ) def _get_default_ckpt(): from online_learner import HF_TOKEN from huggingface_hub import hf_hub_download return hf_hub_download('qoa/kaizen-42m', 'phase4_latest.pt', token=HF_TOKEN) DEFAULT_CKPT = _get_default_ckpt() KAIZEN_STORE = os.path.join(os.path.expanduser('~'), '.kaizen', 'forgetting_memory') def build_domains(): """Split 50 BENCHMARK_TASKS into 5 domains of 10 tasks each. These tasks are naturally diverse (different subjects/types) with high pairwise embedding distances, confirmed clean by Phase 5.""" tasks = list(BENCHMARK_TASKS) domains = [tasks[i*10:(i+1)*10] for i in range(5)] names = [f'D{i+1}({domains[i][0][2][:4]})' for i in range(5)] return names, domains @torch.no_grad() def eval_retention(model, tokenizer, tasks, memory=None, shared_adapter=None): """Mean F1 on tasks (exact same questions used in training — retention test).""" f1_sum = 0.0 for question, answer, _ in tasks: prompt_ids = build_prompt_ids(tokenizer, question)[:BLOCK_SIZE - MAX_GEN] x_prompt = torch.tensor([prompt_ids], dtype=torch.long) task_emb = model.embed_task(x_prompt, adapter=None) adapter = memory.retrieve_merged(task_emb) if memory is not None else shared_adapter gen_ids = generate(model, tokenizer, prompt_ids, adapter=adapter) ref_ids = clean_ids(tokenizer, answer) f1_sum += token_f1(gen_ids, ref_ids) return f1_sum / len(tasks) def kaizen_train_domain(model, tokenizer, memory, tasks): n_stored = 0 for question, answer, _ in tasks: prompt_ids = build_prompt_ids(tokenizer, question)[:BLOCK_SIZE - MAX_GEN] x_prompt = torch.tensor([prompt_ids], dtype=torch.long) ref_ids = clean_ids(tokenizer, answer) with torch.no_grad(): task_emb = model.embed_task(x_prompt, adapter=None) existing = memory.retrieve_merged(task_emb) with torch.no_grad(): gen_ids = generate(model, tokenizer, prompt_ids, adapter=existing) if token_f1(gen_ids, ref_ids) >= ATTEMPT_THRESHOLD: continue new_adp = LoRAAdapter(model.N_LAYERS, model.D_MODEL, LORA_RANK, LORA_ALPHA) if existing: new_adp.load_state_dict(existing.state_dict()) x_upd, y_upd = build_update_seq(tokenizer, question, answer) online_update(model, new_adp, x_upd, y_upd) with torch.no_grad(): post_gen = generate(model, tokenizer, prompt_ids, adapter=new_adp) if token_f1(post_gen, ref_ids) >= STORE_THRESHOLD: memory.add(task_emb, new_adp, {'q': question[:60]}) n_stored += 1 return n_stored def shared_train_domain(model, tokenizer, shared_adapter, tasks): for question, answer, _ in tasks: prompt_ids = build_prompt_ids(tokenizer, question)[:BLOCK_SIZE - MAX_GEN] ref_ids = clean_ids(tokenizer, answer) with torch.no_grad(): gen_ids = generate(model, tokenizer, prompt_ids, adapter=shared_adapter) if token_f1(gen_ids, ref_ids) >= ATTEMPT_THRESHOLD: continue x_upd, y_upd = build_update_seq(tokenizer, question, answer) online_update(model, shared_adapter, x_upd, y_upd) def cl_metrics(R, N): """Standard CL metrics (López-Paz 2017). R[j][i] = retention F1 on domain i after training domain j (0-indexed). """ aa = sum(R[N-1][i] for i in range(N)) / N bwt = sum(R[N-1][i] - R[i][i] for i in range(N-1)) / max(N-1,1) fgt = sum(max(R[j][i] for j in range(i, N)) - R[N-1][i] for i in range(N-1)) / max(N-1,1) return aa, bwt, fgt def main(): t0 = time.time() print('Phase 7.4 — Zero-forgetting benchmark') print('=' * 64) print('Domains: 5 × 10 tasks from BENCHMARK_TASKS (validated diverse set)') print('Test : SAME training questions (retention, not generalization)') tok_file = hf_hub_download(HF_TOK_REPO, 'tokenizer.json', token=HF_TOKEN, cache_dir=None) tokenizer = Tokenizer.from_file(tok_file) model = KaizenWithLoRA() model.load_base(DEFAULT_CKPT) model.eval() print(f'Model loaded ({time.time()-t0:.0f}s)') names, domains = build_domains() N = len(names) print(f'\n{N} domains × 10 tasks = 50 total: {names}') for i, (name, tasks) in enumerate(zip(names, domains)): print(f' D{i+1}: {tasks[0][0][:50]}...') print() # Sanity check: verify cross-domain embedding distances > DIST_THRESHOLD print('Checking inter-domain embedding distances (expect all > 5.0)...') domain_embs = [] for tasks in domains: q = tasks[0][0] prompt_ids = build_prompt_ids(tokenizer, q)[:BLOCK_SIZE - MAX_GEN] x = torch.tensor([prompt_ids], dtype=torch.long) with torch.no_grad(): domain_embs.append(model.embed_task(x, adapter=None)) min_cross_dist = float('inf') for i in range(N): for j in range(N): if i != j: d = ((domain_embs[i] - domain_embs[j])**2).sum().item() min_cross_dist = min(min_cross_dist, d) print(f' min cross-domain squared-L2 dist = {min_cross_dist:.2f} ' f'(DIST_THRESHOLD=5.0, need >> 5.0)') if min_cross_dist < 5.0: print(' WARNING: some domains are too close — retrieval may contaminate!') else: print(' OK: all cross-domain pairs well-separated') print() # ── KAIZEN ──────────────────────────────────────────────────────────────── print('── KAIZEN (per-task episodic adapters) ──') memory = TaskMemory(KAIZEN_STORE, top_k=TOP_K) R_k = [[0.0]*N for _ in range(N)] for j, (dname, tasks) in enumerate(zip(names, domains)): n_stored = kaizen_train_domain(model, tokenizer, memory, tasks) for i in range(j+1): R_k[j][i] = eval_retention(model, tokenizer, domains[i], memory=memory) row = ' '.join(f'{names[i][:6]}={R_k[j][i]:.3f}' for i in range(j+1)) print(f' d{j+1} {dname}: +{n_stored} stored, mem={len(memory)} | {row}') kaizen_aa, kaizen_bwt, kaizen_fgt = cl_metrics(R_k, N) print(f' → AA={kaizen_aa:.4f} BWT={kaizen_bwt:+.4f} Fgt={kaizen_fgt:.4f} ' f'mem={len(memory)}\n') # ── SHARED ──────────────────────────────────────────────────────────────── print('── SHARED single adapter (naive forgetting baseline) ──') shared = LoRAAdapter(model.N_LAYERS, model.D_MODEL, LORA_RANK, LORA_ALPHA) R_s = [[0.0]*N for _ in range(N)] for j, (dname, tasks) in enumerate(zip(names, domains)): shared_train_domain(model, tokenizer, shared, tasks) for i in range(j+1): R_s[j][i] = eval_retention(model, tokenizer, domains[i], shared_adapter=shared) row = ' '.join(f'{names[i][:6]}={R_s[j][i]:.3f}' for i in range(j+1)) print(f' d{j+1} {dname}: | {row}') shared_aa, shared_bwt, shared_fgt = cl_metrics(R_s, N) print(f' → AA={shared_aa:.4f} BWT={shared_bwt:+.4f} Fgt={shared_fgt:.4f}\n') # ── Results ─────────────────────────────────────────────────────────────── print('=' * 64) print(f'{"":24s} {"AA":>8s} {"BWT":>8s} {"Fgt":>8s}') print(f'{"KAIZEN episodic":24s} {kaizen_aa:8.4f} {kaizen_bwt:+8.4f} {kaizen_fgt:8.4f}') print(f'{"SHARED adapter":24s} {shared_aa:8.4f} {shared_bwt:+8.4f} {shared_fgt:8.4f}') print() confirmed = kaizen_bwt > -0.05 and shared_bwt < -0.05 print(f'ZERO-FORGETTING-BY-CONSTRUCTION: {"CONFIRMED ✓" if confirmed else "CHECK MANUALLY"}') print(f' KAIZEN BWT = {kaizen_bwt:+.4f} (≈ 0 expected)') print(f' SHARED BWT = {shared_bwt:+.4f} (negative = forgetting)') print(f'\nRuntime: {time.time()-t0:.0f}s') if __name__ == '__main__': main()