Part E: Adaptation Track

# examples/cap_adapt_starter.py
"""Adaptation track starter: does fine-tuning beat the prompted baseline?

Task: ticket category (6 classes). Model: TinyLM (1.07M parameters).
  1. Prompted baseline: after a prompt that lists the categories, score every
     label by its mean log-probability per token and pick the best (no training).
  2. Choose a recipe WITHOUT touching test: split the 48 dev tickets into 32 for
     training and 16 for validation, try each recipe, keep the best validation
     accuracy among recipes whose general-capability regression stays in budget.
  3. Retrain the chosen recipe on all 48 dev tickets, score the 24 test tickets
     once, and compare with the prompted and keyword baselines (paired sign test).
  4. General-capability regression: perplexity on held-out support dialogue the
     base model was pretrained on, before and after.
Full fine-tuning (every weight), loss on the label tokens only. Saves to
models/cap-triage/<recipe>, never over models/tinylm-base. About 1 to 2 minutes on 1 CPU thread.
"""
from __future__ import annotations

import sys
import time
from pathlib import Path

import torch

from supportdesk.data import CATEGORIES, DATA_DIR, Ticket, load_tickets
from supportdesk.tinylm import batches, load, loss_on, perplexity, save

sys.path.insert(0, str(Path(__file__).resolve().parent))
from cap_baseline import keyword_triage, sign_test_p, wilson  # noqa: E402

torch.set_num_threads(1)
OUT_DIR = Path("models/cap-triage")
HEADER = "Categories: billing, cancellation, account_access, bug, how_to, feature_request.\n"
STEPS, BATCH = 150, 8
RECIPES = {"no_replay": {"lr": 1e-3, "replay": 0.0}, "replay": {"lr": 1e-3, "replay": 1.0}}
MAX_PPL_INCREASE = 0.10  # regression budget: held-out perplexity may rise at most 10%


def prompt_for(text: str, tok, room: int) -> list[int]:
    """Header + ticket (cut from the end if needed) + 'Category:' within `room` tokens."""
    head, tail = tok.encode(HEADER + "Ticket: ").ids, tok.encode("\nCategory:").ids
    body = tok.encode(text.replace("\n", " ")).ids[: room - len(head) - len(tail)]
    return head + body + tail


def label_ids(tok) -> dict[str, list[int]]:
    return {c: tok.encode(" " + c + "\n").ids for c in CATEGORIES}


@torch.no_grad()
def label_scores(model, tok, text: str) -> dict[str, float]:
    """Mean log-probability per token of each label after the prompt.

    Dividing by the label length matters: summed log-probabilities favour short
    labels, and prompted TinyLM then answers 'billing' (2 tokens) for every ticket.
    """
    labels = label_ids(tok)
    room = model.cfg.context - max(len(v) for v in labels.values())
    p = prompt_for(text, tok, room)
    scores = {}
    for cat, ids in labels.items():
        logp = torch.log_softmax(model(torch.tensor([p + ids])[:, :-1])[0], dim=-1)
        scores[cat] = sum(logp[len(p) - 1 + i, t].item() for i, t in enumerate(ids)) / len(ids)
    return scores


def classify(model, tok, text: str) -> str:
    scores = label_scores(model, tok, text)
    return max(scores, key=scores.get)


def predictions(model, tok, tickets: list[Ticket]) -> list[str]:
    model.eval()
    return [classify(model, tok, t.text) for t in tickets]


def training_batch(model, tok, tickets, gen: torch.Generator):
    """Random batch of (x, y, mask); the mask keeps loss on the label tokens only."""
    labels = label_ids(tok)
    room = model.cfg.context - max(len(v) for v in labels.values())
    pick = torch.randint(0, len(tickets), (BATCH,), generator=gen).tolist()
    seqs = [(prompt_for(tickets[i].text, tok, room), labels[tickets[i].gold["category"]]) for i in pick]
    width = max(len(p) + len(lab) for p, lab in seqs) - 1
    x = torch.zeros(BATCH, width, dtype=torch.long)
    y = torch.zeros(BATCH, width, dtype=torch.long)
    m = torch.zeros(BATCH, width)
    for row, (p, lab) in enumerate(seqs):
        ids = p + lab
        x[row, : len(ids) - 1] = torch.tensor(ids[:-1])
        y[row, : len(ids) - 1] = torch.tensor(ids[1:])
        m[row, len(p) - 1: len(ids) - 1] = 1.0
    return x, y, m


def corpus_split() -> tuple[str, str]:
    """Pretraining corpus as (train part, held-out 6,000 characters from the last 5%)."""
    text = (DATA_DIR / "corpus.txt").read_text(encoding="utf-8")
    cut = int(len(text) * 0.95)
    return text[:cut], text[cut:][:6000]


def finetune(train: list[Ticket], lr: float, replay: float, replay_ids: torch.Tensor):
    """Full fine-tuning from the base model; replay adds `replay` x the loss on pretraining text."""
    torch.manual_seed(0)
    model, tok = load()
    gen = torch.Generator().manual_seed(0)
    replay_iter = batches(replay_ids, model.cfg.context, 4, torch.Generator().manual_seed(1))
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.0)
    model.train()
    for _ in range(STEPS):
        x, y, m = training_batch(model, tok, train, gen)
        loss = loss_on(model, x, y, m)
        if replay:
            rx, ry = next(replay_iter)
            loss = loss + replay * loss_on(model, rx, ry)
        opt.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
    model.eval()
    return model, tok


def correct(preds: list[str], tickets: list[Ticket]) -> list[bool]:
    return [p == t.gold["category"] for p, t in zip(preds, tickets)]


def main() -> None:
    started = time.perf_counter()
    dev, test = load_tickets("dev"), load_tickets("test")
    fit, val = [t for i, t in enumerate(dev) if i % 3], [t for i, t in enumerate(dev) if i % 3 == 0]
    corpus_train, heldout = corpus_split()
    base, tok = load()
    replay_ids = torch.tensor(tok.encode(corpus_train).ids)
    ppl0 = perplexity(base, tok, heldout)
    n = len(test)

    kw = [keyword_triage(t.text)[0] == t.gold["category"] for t in test]
    prompted_preds = predictions(base, tok, test)
    prompted = correct(prompted_preds, test)
    print(f"keyword baseline   test {sum(kw):2}/{n}  95% CI {wilson(sum(kw), n)}")
    print(f"prompted TinyLM    test {sum(prompted):2}/{n}  95% CI {wilson(sum(prompted), n)}  "
          f"held-out ppl {ppl0:.3f}  (labels predicted: {sorted(set(prompted_preds))})")

    print(f"\nrecipe selection: train on {len(fit)} dev tickets, validate on {len(val)} (test untouched)")
    best = None
    for name, cfg in RECIPES.items():
        model, _ = finetune(fit, cfg["lr"], cfg["replay"], replay_ids)
        k = sum(correct(predictions(model, tok, val), val))
        ppl = perplexity(model, tok, heldout)
        in_budget = ppl / ppl0 - 1 <= MAX_PPL_INCREASE
        print(f"   {name:10} lr {cfg['lr']} replay {cfg['replay']}: val {k:2}/{len(val)}  "
              f"ppl {ppl0:.3f} -> {ppl:.3f} ({(ppl / ppl0 - 1) * 100:+.0f}%)  "
              f"{'within' if in_budget else 'OVER'} regression budget")
        if in_budget and (best is None or k > best[1]):
            best = (name, k)
    if best is None:
        print("no recipe stayed within the regression budget; stop and rethink the data or the recipe")
        return

    name, cfg = best[0], RECIPES[best[0]]
    print(f"\nchosen recipe: {name}; retrain on all {len(dev)} dev tickets, then score test once")
    model, _ = finetune(dev, cfg["lr"], cfg["replay"], replay_ids)
    tuned = correct(predictions(model, tok, test), test)
    ppl1 = perplexity(model, tok, heldout)
    print(f"fine-tuned TinyLM  test {sum(tuned):2}/{n}  95% CI {wilson(sum(tuned), n)}  "
          f"held-out ppl {ppl0:.3f} -> {ppl1:.3f} ({(ppl1 / ppl0 - 1) * 100:+.0f}%)")
    for label, other in (("prompted TinyLM", prompted), ("keyword baseline", kw)):
        wins = sum(a and not b for a, b in zip(tuned, other))
        losses = sum(b and not a for a, b in zip(tuned, other))
        print(f"   vs {label:16} wins {wins:2} losses {losses:2}  one-sided sign test p = {sign_test_p(wins, losses):.3f}")
    save(model, tok, OUT_DIR / name, meta={"task": "category", "steps": STEPS, "batch": BATCH, **cfg,
                                           "train": "dev split (48)", "base": "tinylm-base"})
    print(f"saved to {OUT_DIR / name}; wall time {time.perf_counter() - started:.0f} s")


if __name__ == "__main__":
    main()