"""Supervised fine-tune of TinyLiquid for forensic analysis (SOP + scratchpad). Format per example: <|persona|><|user|>USER<|assistant|>ASSISTANT<|endoftext|> Loss is masked to the ASSISTANT segment (including scratchpad markers). Usage: .venv/bin/python train/train_sft.py --base ckpt/nlp --data data/sft_forensic.jsonl \ --ckpt ckpt/forensic --epochs 3 """ import argparse import json import math import random import time from pathlib import Path import torch import torch.nn.functional as F from model.config import TinyLiquidConfig, CONFIGS from model.utils import latest_ckpt from model.tiny_liquid import TinyLiquid from data.tokenizer import load_tokenizer USER_T = "<|user|>" ASST_T = "<|assistant|>" EOT_T = "<|endoftext|>" PERSONA_T = {"analyst": "<|analyst|>", "skeptic": "<|skeptic|>", "spock": "<|analyst|>"} def parse_args(): ap = argparse.ArgumentParser() ap.add_argument("--base", default="ckpt/nlp", help="dir with pretrain checkpoints") ap.add_argument("--data", default="data/sft_forensic.jsonl") ap.add_argument("--tok", default="data/tokenizer.json") ap.add_argument("--ckpt", default="ckpt/forensic") ap.add_argument("--epochs", type=int, default=3) ap.add_argument("--batch", type=int, default=8) ap.add_argument("--seq", type=int, default=256) ap.add_argument("--lr", type=float, default=5e-5) ap.add_argument("--eval-every", type=int, default=200) ap.add_argument("--log-every", type=int, default=25) ap.add_argument("--seed", type=int, default=7) ap.add_argument("--threads", type=int, default=8) return ap.parse_args() def tokenize_example(tok, ex, seq, u_id, a_id, eot_id, p_ids): persona = PERSONA_T.get(ex["persona"], PERSONA_T["analyst"]) p_id = p_ids[ex["persona"]] if ex["persona"] in p_ids else p_ids["analyst"] parts = [tok.encode(persona).ids, [u_id], tok.encode(ex["user"]).ids, [a_id], tok.encode(ex["assistant"]).ids, [eot_id]] ids = [i for part in parts for i in part] if len(ids) > seq: # truncate assistant side keep = seq - 1 ids = ids[:keep] + [eot_id] asst_start = len(parts[0]) + 1 + len(parts[2]) + 1 # index of first assistant token x = torch.tensor(ids[:-1], dtype=torch.long) y = torch.tensor(ids[1:], dtype=torch.long) mask = torch.zeros_like(y, dtype=torch.bool) mask[asst_start - 1:] = True # y positions that predict assistant tokens mask = mask[: x.shape[0]] return x, y[: x.shape[0]], mask, p_id def collate(items, seq): xs, ys, ms, ps = [], [], [], [] for x, y, m, p in items: xs.append(F.pad(x, (0, seq - x.shape[0]), value=0)) ys.append(F.pad(y, (0, seq - y.shape[0]), value=0)) ms.append(F.pad(m, (0, seq - m.shape[0]), value=False)) ps.append(p) return (torch.stack(xs), torch.stack(ys), torch.stack(ms), torch.tensor(ps, dtype=torch.long)) def main(): args = parse_args() torch.set_num_threads(args.threads) torch.manual_seed(args.seed) random.seed(args.seed) rng = random.Random(args.seed) tok = load_tokenizer(args.tok) u_id = tok.token_to_id(USER_T) a_id = tok.token_to_id(ASST_T) eot_id = tok.token_to_id(EOT_T) # persona embedding indices (NOT tokenizer ids): 0=none, 1=analyst, 2=skeptic p_ids = {"analyst": 1, "skeptic": 2} assert None not in (u_id, a_id, eot_id), "special tokens missing from tokenizer" examples = [json.loads(l) for l in open(args.data, encoding="utf-8") if l.strip()] rng.shuffle(examples) n_eval = min(128, len(examples) // 10) eval_ex, train_ex = examples[:n_eval], examples[n_eval:] print(f"train {len(train_ex)} eval {len(eval_ex)}", flush=True) base_ckpt = latest_ckpt(args.base) assert base_ckpt, f"no pretrain checkpoint in {args.base}" base = torch.load(base_ckpt, map_location="cpu") config = base.get("config") or CONFIGS["tiny10m"] cfg = TinyLiquidConfig(vocab_size=tok.get_vocab_size(), **{k: v for k, v in config.items() if k != "vocab_size"}) model = TinyLiquid(cfg) model.load_state_dict(base["model"]) opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.05) print(f"loaded base {base_ckpt.name} (step {base.get('step', '?')})", flush=True) out_dir = Path(args.ckpt) out_dir.mkdir(parents=True, exist_ok=True) steps_per_epoch = max(1, len(train_ex) // args.batch) total_steps = steps_per_epoch * args.epochs model.train() def make_items(exs): return [tokenize_example(tok, e, args.seq, u_id, a_id, eot_id, p_ids) for e in exs] def run_eval(): model.eval() items = make_items(eval_ex) total, n = 0.0, 0 for i in range(0, len(items), args.batch): x, y, m, p = collate(items[i:i + args.batch], args.seq) with torch.no_grad(): logits = model(x, persona_ids=p) logits = logits.reshape(-1, logits.size(-1)) loss = F.cross_entropy(logits, y.reshape(-1), reduction="none") loss = (loss * m.reshape(-1)).sum() / m.sum() total += loss.item() * m.sum().item() n += m.sum().item() model.train() return total / n t0 = time.time() step = 0 for ep in range(args.epochs): rng.shuffle(train_ex) items = make_items(train_ex) for i in range(0, len(items) - len(items) % args.batch, args.batch): step += 1 x, y, m, p = collate(items[i:i + args.batch], args.seq) opt.zero_grad(set_to_none=True) logits = model(x, persona_ids=p) logits = logits.reshape(-1, logits.size(-1)) loss = F.cross_entropy(logits, y.reshape(-1), reduction="none") loss = (loss * m.reshape(-1)).sum() / m.sum() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() if step % args.log_every == 0: dt = time.time() - t0 print(f"step {step}/{total_steps} loss {loss.item():.4f} " f"{args.batch*args.seq*args.log_every/dt:.0f} tok/s", flush=True) t0 = time.time() if step % args.eval_every == 0: vl = run_eval() print(f" [eval {step}] sft_val_loss {vl:.4f}", flush=True) torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": step, "config": cfg.__dict__}, str(out_dir / f"model_{step}.pt")) torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": step, "config": cfg.__dict__}, str(out_dir / "model_final.pt")) print(f"done -> {out_dir}", flush=True) if __name__ == "__main__": main()