train/train_sft.py

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