"""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()