forgetting_benchmark.py

#!/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()