scale_experiment.py

#!/usr/bin/env python3
"""
scale_experiment.py — Scale comparison: KAIZEN 42M vs GPT-2 117M.

Research question: does semantic structure recoverable by a linear head
from frozen LM hidden states improve with model scale?

Protocol (identical to Phase 7.0 / semantic_probe.py):
  - Same 50 paraphrase pairs from probe_paraphrase.json
  - Same 70/30 task-level split (35 train, 15 held-out)
  - Same LinearHead architecture (d_model→128, InfoNCE)
  - Same recall@1 metric on held-out tasks

Output:
  - Raw recall@1 (no head) at each scale = lower bound
  - Recall@1 learning curve across epochs [0,10,25,50,100,150,200,300]
  - Final recall@1 = publishable single-number comparison

KAIZEN 42M: embed_task via canonical prompt [BOS,USER]+q+[ASST], mean-pool, 512-dim
GPT-2 117M: plain question text tokenized with GPT-2 BPE, mean-pool last hidden, 768-dim
"""
import json, random, sys, time
import torch
import torch.nn as nn
import torch.nn.functional as F
from tokenizers import Tokenizer
from transformers import GPT2Model, GPT2Tokenizer
from huggingface_hub import hf_hub_download

from lora import KaizenWithLoRA
from eval_benchmark import BENCHMARK_TASKS, build_prompt_ids, BLOCK_SIZE, MAX_GEN
from online_learner import HF_TOKEN, HF_TOK_REPO

# ── Config ────────────────────────────────────────────────────────────────────
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()
PROBE_PATH  = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'probe_paraphrase.json')
SEED        = 42
TRAIN_FRAC  = 0.70
HEAD_DIM    = 128
HEAD_LR     = 1e-3
HEAD_TEMP   = 0.07
CHECKPOINTS = [0, 10, 25, 50, 100, 150, 200, 300]   # epochs at which to record recall@1


# ── LinearHead + InfoNCE ──────────────────────────────────────────────────────
class LinearHead(nn.Module):
    def __init__(self, d_in: int, d_out: int = HEAD_DIM):
        super().__init__()
        self.proj = nn.Linear(d_in, d_out, bias=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.normalize(self.proj(x), dim=-1)


def infonce_loss(a: torch.Tensor, b: torch.Tensor,
                 temp: float = HEAD_TEMP) -> torch.Tensor:
    """Bidirectional InfoNCE on (anchor, positive) pairs. a, b: [B, D]."""
    sim    = a @ b.T / temp               # [B, B]
    labels = torch.arange(sim.shape[0])
    return (F.cross_entropy(sim, labels) + F.cross_entropy(sim.T, labels)) / 2


# ── Metrics ───────────────────────────────────────────────────────────────────
def recall_at_1_head(head: LinearHead,
                     anc_embs: torch.Tensor,
                     pos_embs: torch.Tensor,
                     held_out_idx: list) -> float:
    """Recall@1 in g-space.

    For each held-out anchor i: query = head(anc_embs[i]).
    Gallery = head(pos_embs) over ALL 50 tasks (train + held-out).
    Correct if nearest gallery entry == i (the true paraphrase).
    """
    with torch.no_grad():
        g_pos = head(pos_embs)                       # [50, 128]
        g_anc = head(anc_embs[held_out_idx])         # [N_held, 128]

    correct = 0
    for k, i in enumerate(held_out_idx):
        dists   = ((g_pos - g_anc[k].unsqueeze(0)) ** 2).sum(dim=1)  # [50]
        nearest = int(dists.argmin().item())
        if nearest == i:
            correct += 1
    return correct / len(held_out_idx)


def recall_at_1_raw(anc_embs: torch.Tensor,
                    pos_embs: torch.Tensor,
                    held_out_idx: list) -> float:
    """Recall@1 in raw embedding space (L2-normalized, cosine-equivalent)."""
    with torch.no_grad():
        a_n = F.normalize(anc_embs, dim=-1)          # [50, d]
        p_n = F.normalize(pos_embs, dim=-1)          # [50, d]
    correct = 0
    for i in held_out_idx:
        dists   = ((p_n - a_n[i].unsqueeze(0)) ** 2).sum(dim=1)
        nearest = int(dists.argmin().item())
        if nearest == i:
            correct += 1
    return correct / len(held_out_idx)


# ── Training loop ─────────────────────────────────────────────────────────────
def train_with_checkpoints(d_in: int,
                           anc_train: torch.Tensor,
                           pos_train: torch.Tensor,
                           anc_all: torch.Tensor,
                           pos_all: torch.Tensor,
                           held_out_idx: list) -> dict:
    """Train LinearHead; return {epoch: recall@1} at each CHECKPOINT epoch."""
    head      = LinearHead(d_in=d_in)
    optimizer = torch.optim.Adam(head.parameters(), lr=HEAD_LR)
    N         = anc_train.shape[0]
    results   = {}

    if 0 in CHECKPOINTS:
        results[0] = recall_at_1_head(head, anc_all, pos_all, held_out_idx)

    max_ep = max(CHECKPOINTS)
    for ep in range(1, max_ep + 1):
        perm = torch.randperm(N)
        optimizer.zero_grad()
        loss = infonce_loss(head(anc_train[perm]), head(pos_train[perm]))
        loss.backward()
        optimizer.step()
        if ep in CHECKPOINTS:
            results[ep] = recall_at_1_head(head, anc_all, pos_all, held_out_idx)

    return results


# ── Embedding extractors ──────────────────────────────────────────────────────
def embed_kaizen(model: KaizenWithLoRA, tokenizer,
                 questions: list) -> torch.Tensor:
    """Canonical prompt embed_task for all questions. Returns [N, 512]."""
    vecs = []
    with torch.no_grad():
        for q in questions:
            pids = build_prompt_ids(tokenizer, q)[:BLOCK_SIZE - MAX_GEN]
            x    = torch.tensor([pids], dtype=torch.long)
            vecs.append(model.embed_task(x, adapter=None))   # [512]
    return torch.stack(vecs)


def embed_gpt2(model: GPT2Model, tokenizer: GPT2Tokenizer,
               questions: list) -> torch.Tensor:
    """Mean-pool GPT-2 last hidden state for all questions. Returns [N, 768]."""
    vecs = []
    with torch.no_grad():
        for q in questions:
            inputs = tokenizer(q, return_tensors='pt',
                               truncation=True, max_length=512)
            out    = model(**inputs)
            # last_hidden_state: [1, seq_len, 768] → mean over seq_len
            vecs.append(out.last_hidden_state[0].mean(dim=0))   # [768]
    return torch.stack(vecs)


# ── Main ──────────────────────────────────────────────────────────────────────
def main():
    t0  = time.time()
    rng = random.Random(SEED)

    print('Scale experiment: KAIZEN 42M vs GPT-2 117M')
    print('Metric: recall@1 on held-out paraphrase pairs')
    print(f'Checkpoints: {CHECKPOINTS}')
    print('=' * 64)

    # ── Probe data ───────────────────────────────────────────────────────────
    with open(PROBE_PATH) as f:
        paraphrases = json.load(f)   # list of [para_q, answer, category]

    orig_qs = [t[0] for t in BENCHMARK_TASKS]   # 50 original questions
    para_qs = [p[0] for p in paraphrases]        # 50 paraphrase questions

    # 70/30 task-level split — same seed as Phase 7.0
    all_idx = list(range(50))
    rng.shuffle(all_idx)
    n_train      = int(TRAIN_FRAC * 50)   # 35
    train_idx    = all_idx[:n_train]
    held_out_idx = all_idx[n_train:]      # 15
    print(f'Train tasks: {n_train}  |  Held-out tasks: {len(held_out_idx)}')

    # ── KAIZEN 42M embeddings ─────────────────────────────────────────────────
    print('\n[1/4] Loading KAIZEN 42M and embedding probe questions...')
    tok_file = hf_hub_download(HF_TOK_REPO, 'tokenizer.json',
                                token=HF_TOKEN, cache_dir=None)
    kz_tok   = Tokenizer.from_file(tok_file)
    kz_model = KaizenWithLoRA()
    kz_model.load_base(DEFAULT_CKPT)
    kz_model.eval()
    kz_params = sum(p.numel() for p in kz_model.parameters())
    print(f'  KAIZEN: {kz_params/1e6:.1f}M params')

    kz_anc = embed_kaizen(kz_model, kz_tok, orig_qs)   # [50, 512]
    kz_pos = embed_kaizen(kz_model, kz_tok, para_qs)   # [50, 512]
    del kz_model   # free model weights; tensors remain valid
    print(f'  Embeddings: {list(kz_anc.shape)}  (elapsed {time.time()-t0:.0f}s)')

    # ── GPT-2 117M embeddings ─────────────────────────────────────────────────
    print('\n[2/4] Loading GPT-2 117M and embedding probe questions...')
    gpt2_tok   = GPT2Tokenizer.from_pretrained('gpt2',
                                                cache_dir='/tmp/gpt2_cache')
    gpt2_model = GPT2Model.from_pretrained('gpt2',
                                            cache_dir='/tmp/gpt2_cache')
    gpt2_model.eval()
    gpt2_params = sum(p.numel() for p in gpt2_model.parameters())
    print(f'  GPT-2: {gpt2_params/1e6:.1f}M params')

    g2_anc = embed_gpt2(gpt2_model, gpt2_tok, orig_qs)   # [50, 768]
    g2_pos = embed_gpt2(gpt2_model, gpt2_tok, para_qs)   # [50, 768]
    del gpt2_model   # free model weights
    print(f'  Embeddings: {list(g2_anc.shape)}  (elapsed {time.time()-t0:.0f}s)')

    # ── Raw recall@1 (no head) ────────────────────────────────────────────────
    kz_raw = recall_at_1_raw(kz_anc, kz_pos, held_out_idx)
    g2_raw = recall_at_1_raw(g2_anc, g2_pos, held_out_idx)
    print(f'\nRaw recall@1 (frozen embeddings, no head, cosine NN):')
    print(f'  KAIZEN 42M  : {kz_raw:.4f}')
    print(f'  GPT-2  117M : {g2_raw:.4f}')
    print(f'  Δ (117M−42M): {g2_raw - kz_raw:+.4f}')

    # ── Head training: build train tensors ────────────────────────────────────
    kz_anc_tr = kz_anc[train_idx]   # [35, 512]
    kz_pos_tr = kz_pos[train_idx]   # [35, 512]
    g2_anc_tr = g2_anc[train_idx]   # [35, 768]
    g2_pos_tr = g2_pos[train_idx]   # [35, 768]

    # ── Train KAIZEN head ─────────────────────────────────────────────────────
    print(f'\n[3/4] Training LinearHead for KAIZEN 42M (512→{HEAD_DIM}, InfoNCE)...')
    kz_curve = train_with_checkpoints(
        d_in=512,
        anc_train=kz_anc_tr, pos_train=kz_pos_tr,
        anc_all=kz_anc, pos_all=kz_pos,
        held_out_idx=held_out_idx,
    )
    print(f'  Done (elapsed {time.time()-t0:.0f}s)')

    # ── Train GPT-2 head ──────────────────────────────────────────────────────
    print(f'[4/4] Training LinearHead for GPT-2 117M (768→{HEAD_DIM}, InfoNCE)...')
    g2_curve = train_with_checkpoints(
        d_in=768,
        anc_train=g2_anc_tr, pos_train=g2_pos_tr,
        anc_all=g2_anc, pos_all=g2_pos,
        held_out_idx=held_out_idx,
    )
    print(f'  Done (elapsed {time.time()-t0:.0f}s)')

    # ── Results table ─────────────────────────────────────────────────────────
    print()
    print('=' * 64)
    print('Recall@1 on held-out paraphrase pairs (linear head, InfoNCE, 35 train)')
    print(f'{"Epoch":>8s}  {"KAIZEN 42M":>12s}  {"GPT-2 117M":>12s}  {"Δ(117M-42M)":>12s}')
    print('-' * 52)
    print(f'{"raw":>8s}  {kz_raw:12.4f}  {g2_raw:12.4f}  {g2_raw-kz_raw:+12.4f}')
    for ep in CHECKPOINTS:
        kz_r = kz_curve.get(ep, float('nan'))
        g2_r = g2_curve.get(ep, float('nan'))
        delta = g2_r - kz_r if (kz_r == kz_r and g2_r == g2_r) else float('nan')
        print(f'{ep:8d}  {kz_r:12.4f}  {g2_r:12.4f}  {delta:+12.4f}')

    kz_final = kz_curve.get(max(CHECKPOINTS), 0.0)
    g2_final = g2_curve.get(max(CHECKPOINTS), 0.0)

    print()
    print(f'Summary (epoch {max(CHECKPOINTS)}):')
    print(f'  KAIZEN 42M  raw={kz_raw:.4f} → head={kz_final:.4f}  '
          f'(gain={kz_final-kz_raw:+.4f})')
    print(f'  GPT-2  117M raw={g2_raw:.4f} → head={g2_final:.4f}  '
          f'(gain={g2_final-g2_raw:+.4f})')
    print(f'  Scale benefit at epoch 300: {g2_final - kz_final:+.4f} '
          f'({"117M > 42M" if g2_final > kz_final else "42M >= 117M"})')
    print(f'\nRuntime: {time.time()-t0:.0f}s')


if __name__ == '__main__':
    main()