train/train_classifier.py

"""Claim-judgment classifier: verdict / confidence / fallacy from a claim.

Trains a small linear head on the TinyLiquid base's final hidden state. The
base is frozen except the last block + head (light adapter), so the 7.8M model
becomes a reliable claim-conditioned judge instead of a drifting generator.

Usage:
  .venv/bin/python train/train_classifier.py --base ckpt/v8_lora/best.pt \
      --data data/sft_forensic.jsonl --ckpt ckpt/judge
"""
import argparse, json, random, re, time
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from model.config import TinyLiquidConfig
from model.tiny_liquid import TinyLiquid
from data.tokenizer import load_tokenizer

VERDICT_ORDER = ["true statement", "false statement", "supports", "refutes", "not_enough_info"]
CONF_ORDER = ["high", "medium", "low"]
P_IDS = {"analyst": 1, "skeptic": 2, "none": 0}

def parse_label(text, key):
    m = re.search(key + r"\s*:\s*([^.]+)\.", text, re.I)
    if not m:
        return None
    lab = m.group(1).strip().lower()
    if key.lower() == "verdict":
        for cand in VERDICT_ORDER:
            if lab.startswith(cand) or cand.startswith(lab.split(" ")[0][:4]):
                return cand
        if lab.startswith("true"): return "true statement"
        if lab.startswith("false"): return "false statement"
        if lab.startswith("not_enough") or lab.startswith("not enough"): return "not_enough_info"
        if lab.startswith("support"): return "supports"
        if lab.startswith("refut"): return "refutes"
        return None
    if key.lower() == "confidence":
        if lab.startswith("high"): return "high"
        if lab.startswith("medium"): return "medium"
        if lab.startswith("low"): return "low"
        return None
    # fallacy: keep as-is (13 classes)
    return lab

def build(args):
    tok = load_tokenizer(args.tok)
    rows = [json.loads(l) for l in open(args.data, encoding="utf-8") if l.strip()]
    items = []
    for r in rows:
        u = r.get("user", "")
        a = r.get("assistant", "")
        if not u or not a:
            continue
        pid = P_IDS.get(r.get("persona", "analyst"), 1)
        v = parse_label(a, "Verdict")
        c = parse_label(a, "Confidence")
        f = parse_label(a, "Fallacy")
        items.append({"ids": tok.encode(u).ids, "pid": pid, "v": v, "c": c, "f": f})
    print(f"rows {len(rows)} usable {len(items)}", flush=True)
    return tok, items

def make_sets(items, key, valid_vals, seed=17):
    rng = random.Random(seed)
    data = [it for it in items if it[key] in valid_vals]
    rng.shuffle(data)
    n_val = max(64, int(len(data) * 0.12))
    return data[n_val:], data[:n_val], valid_vals

def encode_batch(model, items, tok, max_len=192, grad=True):
    xs, ps = [], []
    for it in items:
        ids = it["ids"][:max_len]
        xs.append(ids)
        ps.append(it["pid"])
    L = max(len(x) for x in xs)
    buf = torch.zeros(len(xs), L, dtype=torch.long)
    for i, x in enumerate(xs):
        buf[i, :len(x)] = torch.tensor(x, dtype=torch.long)
    if grad:
        h = model.encode(buf, persona_ids=torch.tensor(ps))
    else:
        with torch.no_grad():
            h = model.encode(buf, persona_ids=torch.tensor(ps))
    mask = torch.arange(L).unsqueeze(0) < torch.tensor([len(x) for x in xs]).unsqueeze(1)  # (n, L)
    h = h * mask.unsqueeze(-1)
    return h.sum(1) / mask.sum(1, keepdim=True)  # masked mean pool (n, d)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", default="ckpt/v8_lora/best.pt")
    ap.add_argument("--data", default="data/sft_forensic.jsonl")
    ap.add_argument("--tok", default="data/tokenizer.json")
    ap.add_argument("--ckpt", default="ckpt/judge")
    ap.add_argument("--epochs", type=int, default=8)
    ap.add_argument("--batch", type=int, default=32)
    ap.add_argument("--lr", type=float, default=3e-4)
    ap.add_argument("--threads", type=int, default=4)
    args = ap.parse_args()
    torch.set_num_threads(args.threads)
    torch.manual_seed(17)
    tok, items = build(args)

    sd = torch.load(args.base, map_location="cpu")
    cfg = TinyLiquidConfig(vocab_size=tok.get_vocab_size(), **{k: v for k, v in sd["config"].items() if k != "vocab_size"})
    model = TinyLiquid(cfg)
    model.load_state_dict(sd["model"])
    for p in model.parameters():
        p.requires_grad = False
    for p in model.blocks[-2:].parameters():  # adapt last 2 blocks + head
        p.requires_grad = True
    for p in model.norm_out.parameters():
        p.requires_grad = True
    model.train()
    d = cfg.d_model

    heads = {}
    for key, order in [("v", VERDICT_ORDER), ("c", CONF_ORDER), ("f", None)]:
        if key == "f":
            vals = sorted({it["f"] for it in items if it["f"]})
        else:
            vals = order
        if not vals:
            continue
        tr, va, vals = make_sets(items, key, vals)
        head = nn.Linear(d, len(vals))
        idx = {v: i for i, v in enumerate(vals)}
        heads[key] = {"head": head, "train": tr, "val": va, "idx": idx, "vals": vals}
        print(f"head {key}: {len(vals)} classes, train {len(tr)} val {len(va)}", flush=True)

    params = [p for p in model.parameters() if p.requires_grad]
    for hd in heads.values():
        params += list(hd["head"].parameters())
    opt = torch.optim.AdamW(params, lr=args.lr, weight_decay=0.01)

    out = Path(args.ckpt); out.mkdir(parents=True, exist_ok=True)
    t0 = time.time()
    for ep in range(args.epochs):
        for key, hd in heads.items():
            rng = random.Random(ep * 101 + 7)
            rng.shuffle(hd["train"])
        # interleave heads per batch
        for i in range(0, max(len(hd["train"]) for hd in heads.values()), args.batch):
            opt.zero_grad(set_to_none=True)
            loss = 0.0
            for key, hd in heads.items():
                batch = hd["train"][i:i + args.batch]
                if not batch:
                    continue
                h = encode_batch(model, batch, tok)
                logits = hd["head"](h)
                target = torch.tensor([hd["idx"][it[key]] for it in batch])
                loss = loss + F.cross_entropy(logits, target)
            if loss == 0:
                continue
            loss.backward()
            torch.nn.utils.clip_grad_norm_(params, 1.0)
            opt.step()
        # eval
        line = []
        for key, hd in heads.items():
            hd["head"].eval()
            with torch.no_grad():
                h = encode_batch(model, hd["val"], tok, grad=False)
                logits = hd["head"](h)
                preds = logits.argmax(-1)
                targets = torch.tensor([hd["idx"][it[key]] for it in hd["val"]])
                acc = (preds == targets).float().mean().item()
            line.append(f"{key}_acc {acc:.3f}")
            hd["head"].train()
        print(f"epoch {ep+1}/{args.epochs} " + " ".join(line) + f" ({time.time()-t0:.0f}s)", flush=True)
        t0 = time.time()

    torch.save({"heads": {k: {"state": hd["head"].state_dict(), "vals": hd["vals"]} for k, hd in heads.items()},
                "config": cfg.__dict__, "base": args.base}, out / "judge.pt")
    print("saved ->", out / "judge.pt", flush=True)

if __name__ == "__main__":
    main()