backup all: 37 files (final)

tests/test_posttrain.py

@@ -0,0 +1,71 @@
+"""Unit tests for train/ties_merge.py and research/rlvr.py.
+
+Run: .venv/bin/python tests/test_posttrain.py
+"""
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+import torch
+
+from train.ties_merge import ties_merge, trim_delta
+from research.rlvr import reward, reward_card
+
+
+def test_trim_keeps_topk():
+    d = torch.tensor([1.0, -0.5, 0.01, 0.001, 0.0005])
+    t = trim_delta(d, keep=0.4)
+    assert t[0] == 1.0 and t[1] == -0.5
+    assert t[2] == 0.0 and t[3] == 0.0
+
+
+def test_ties_sign_consensus():
+    base = {"w": torch.zeros(4)}
+    t1 = {"w": torch.tensor([1.0, 1.0, 1.0, -1.0])}
+    t2 = {"w": torch.tensor([1.0, -1.0, 1.0, -1.0])}
+    out = ties_merge(base, [t1, t2], keep=1.0)
+    # sign agreement at idx 0, 2, 3 -> merge; idx 1 disagrees -> zero
+    assert out["w"][0] == 1.0
+    assert out["w"][2] == 1.0
+    assert out["w"][3] == -1.0
+    assert out["w"][1] == 0.0
+
+
+def test_rlvr_reward_correct_with_citation():
+    r = reward(gold="refutes", policy="refutes", citation="1982",
+               evidence="the deed file states 1982")
+    assert r["verdict"] == 1.0 and r["citation"] == 0.2 and r["total"] == 1.2
+
+
+def test_rlvr_reward_wrong_verdict():
+    r = reward(gold="refutes", policy="supports", citation="1982",
+               evidence="the deed file states 1982")
+    assert r["verdict"] == -1.0 and r["total"] == -0.8
+
+
+def test_rlvr_reward_abstain_is_zero():
+    r = reward(gold="refutes", policy="not enough information",
+               citation="", evidence="the deed file states 1982")
+    assert r["verdict"] == 0.0 and r["total"] == 0.0
+
+
+def test_rlvr_reward_false_citation_penalty():
+    r = reward(gold="supports", policy="supports", citation="1978",
+               evidence="the deed file states 1982")
+    assert r["verdict"] == 1.0 and r["citation"] == -0.2 and r["total"] == 0.8
+
+
+def test_rlvr_card_trace():
+    card = reward_card(gold="refutes", policy="refutes", citation="1982",
+                       evidence="the deed file states 1982", probe="rt05")
+    assert card["probe"] == "rt05" and card["total"] == 1.2
+    for k in ("gold", "policy", "verdict", "citation", "total"):
+        assert k in card
+
+
+if __name__ == "__main__":
+    fns = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
+    for fn in fns:
+        fn()
+        print(f"ok {fn.__name__}")
+    print(f"\n{len(fns)} posttrain tests passed")

tests/test_verify_loop.py

@@ -0,0 +1,131 @@
+"""Standalone unit tests for research/verify_loop.py.
+
+Run: .venv/bin/python tests/test_verify_loop.py
+"""
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from research.verify_loop import plan_checks, run_checks, verify_case
+
+
+def _verified_sources():
+    return [
+        {"source_id": "dot", "url": "https://records.example/dot",
+         "retrieved_at": "2026-08-12T12:00:00Z", "content_sha256": "a" * 64,
+         "independent": True, "retrievable": True,
+         "triage": {"independence": 3, "proximity": 3, "recency": 2, "track": 3, "interest": 3}},
+        {"source_id": "archive", "url": "https://archive.example/bridge",
+         "retrieved_at": "2026-08-12T12:01:00Z", "content_sha256": "b" * 64,
+         "independent": True, "retrievable": True,
+         "triage": {"independence": 2, "proximity": 2, "recency": 2, "track": 2, "interest": 2}},
+    ]
+
+
+def test_plan_checks_extracts_values():
+    checks = plan_checks('Claim: "the bridge opened in 2010" and cost $4.2M at 9:30am.')
+    kinds = [c["kind"] for c in checks]
+    assert "quote" in kinds and "number" in kinds and "time" in kinds
+    assert len(checks) <= 8
+
+
+def test_run_checks_supports():
+    checks = [{"kind": "number", "value": "2010"}]
+    res = run_checks(checks, lambda v: "The DOT file lists the bridge opening year as 2010.")
+    assert res[0]["verdict"] == "supports"
+
+
+def test_run_checks_refutes():
+    checks = [{"kind": "number", "value": "2010"}]
+    res = run_checks(checks, lambda v: "The DOT file lists the bridge opening year as 2012.")
+    assert res[0]["verdict"] == "refutes"
+
+
+def test_run_checks_no_evidence():
+    checks = [{"kind": "quote", "value": "classified"}]
+    res = run_checks(checks, lambda v: "")
+    assert res[0]["verdict"] == "not enough information"
+    assert res[0]["kind"] == "no-evidence"
+
+
+def test_verify_case_rule_refuted_wins():
+    d = verify_case('Claim: "the bridge opened in 2010"',
+                    draft_verdict="true", draft_conf="HIGH",
+                    retrieve=lambda v: "The DOT file lists the bridge opening year as 2012.",
+                    require_source_policy=False)
+    assert d["verdict"] == "false"
+    assert d["confidence"] == "HIGH"
+    assert d["basis"] == "rule-refuted"
+
+
+def test_verify_case_rule_verified():
+    d = verify_case('Claim: "the bridge opened in 2010"',
+                    draft_verdict="true", draft_conf="MEDIUM",
+                    retrieve=lambda v: "The DOT file lists the bridge opening year as 2010.",
+                    require_source_policy=False)
+    assert d["verdict"] == "true"
+    assert d["confidence"] == "HIGH"
+    assert d["basis"] == "rule-verified"
+
+
+def test_verify_case_unresolved_downgrades_high():
+    d = verify_case("Claim: the memo is significant.",
+                    draft_verdict="overclaim", draft_conf="HIGH",
+                    retrieve=lambda v: "", require_source_policy=False)
+    assert d["confidence"] == "MEDIUM"
+    assert d["basis"] == "draft-high-downgraded-unverified"
+
+
+def test_verify_case_unresolved_keeps_low():
+    d = verify_case("Claim: the memo is significant.",
+                    draft_verdict="not enough information", draft_conf="LOW",
+                    retrieve=lambda v: "", require_source_policy=False)
+    assert d["verdict"] == "not enough information"
+    assert d["confidence"] == "LOW"
+    assert d["abstained"]
+
+
+def test_strict_policy_overrides_an_incorrect_model_draft():
+    d = verify_case('Claim: "the bridge opened in 2010"',
+                    draft_verdict="false", draft_conf="HIGH",
+                    retrieve=lambda v: {
+                        "evidence": "The DOT filing lists the bridge opening year as 2010.",
+                        "sources": _verified_sources(), "claim_relation": "supports"},
+                    require_source_policy=True)
+    assert d["verdict"] == "true"
+    assert d["confidence"] == "HIGH"
+    assert d["basis"] == "source-policy-verified"
+    assert set(d["sources"]) == {"dot", "archive"}
+
+
+def test_strict_policy_fails_closed_on_a_single_untraceable_lead():
+    d = verify_case('Claim: "the bridge opened in 2010"',
+                    draft_verdict="true", draft_conf="HIGH",
+                    retrieve=lambda v: {
+                        "evidence": "An anonymous post says the bridge opened in 2010.",
+                        "sources": [{"source_id": "forum", "url": "https://forum.example/post",
+                                     "retrieved_at": "", "content_sha256": "bad",
+                                     "independent": True, "retrievable": True,
+                                     "triage": {"independence": 1, "proximity": 0,
+                                                "recency": 1, "track": 0, "interest": 0}}],
+                        "claim_relation": "supports"},
+                    require_source_policy=True)
+    assert d["verdict"] == "not enough information"
+    assert d["confidence"] == "LOW"
+    assert d["abstained"]
+
+
+def test_verify_case_defaults_to_fail_closed_source_policy():
+    d = verify_case('Claim: "the bridge opened in 2010"',
+                    draft_verdict="true", draft_conf="HIGH",
+                    retrieve=lambda v: "A copied page says the bridge opened in 2010.")
+    assert d["verdict"] == "not enough information"
+    assert d["confidence"] == "LOW"
+    assert d["abstained"]
+
+
+if __name__ == "__main__":
+    fns = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
+    for fn in fns:
+        fn()
+        print(f"PASS {fn.__name__}")
+    print(f"\n{len(fns)} tests passed")

train/chain_dpo_v26_after_sft.sh

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+# Chains LFM2 DPO v26 after SFT v26 completes (one heavy torch job at a time).
+# Polls for ckpt/hybrid50m_v26_lora/model_final.pt, then launches DPO v26.
+# Launch detached: setsid nohup ./train/chain_dpo_v26_after_sft.sh >/dev/null 2>&1 </dev/null & disown
+set -u
+cd "$(dirname "$0")/.."
+LOG=logs/chain_dpo_v26.log
+echo "=== chain start $(date '+%Y-%m-%d %H:%M:%S') ===" >> "$LOG"
+
+for i in $(seq 1 720); do
+  if [ -f "ckpt/hybrid50m_v26_lora/model_final.pt" ]; then
+    echo "=== SFT v26 complete at $(date '+%H:%M:%S'); launching DPO v26 ===" >> "$LOG"
+    setsid nohup ./train/watchdog_dpo_v26.sh >/dev/null 2>&1 </dev/null & disown
+    echo "DPO v26 watchdog launched" >> "$LOG"
+    exit 0
+  fi
+  sleep 60
+done
+echo "=== chain timeout after 12h; SFT v26 not complete ===" >> "$LOG"
+exit 1

train/chain_post_sft_v22.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+# Chains SFT(final) -> DPO -> eval gate, all detached/self-healing.
+# Does NOT merge or release (gate numbers are reported for the owner to decide).
+set -u
+cd "$(dirname "$0")/.."
+export PYTHONPATH=$PWD MALLOC_ARENA_MAX=2
+SFT_CKPT=ckpt/hybrid50m_v22_lora
+DPO_CKPT=ckpt/hybrid50m_v22_dpo
+CL=logs/chain_v22.log
+echo "[chain $(date '+%H:%M:%S')] waiting for SFT final" >> "$CL"
+while [ ! -f "$SFT_CKPT/model_final.pt" ]; do
+  pgrep -f watchdog_lora_v22 >/dev/null || setsid nohup ./train/watchdog_lora_v22.sh >/dev/null 2>&1 </dev/null &
+  sleep 30
+done
+echo "[chain $(date '+%H:%M:%S')] SFT done -> launching DPO" >> "$CL"
+setsid nohup ./train/watchdog_dpo_v22.sh >/dev/null 2>&1 </dev/null &
+while [ ! -f "$DPO_CKPT/dpo_complete.json" ]; do
+  if [ -f "$DPO_CKPT/dpo_failed.json" ]; then
+    echo "[chain $(date '+%H:%M:%S')] DPO failed; eval is blocked" >> "$CL"
+    exit 3
+  fi
+  pgrep -f watchdog_dpo_v22 >/dev/null || setsid nohup ./train/watchdog_dpo_v22.sh >/dev/null 2>&1 </dev/null &
+  sleep 30
+done
+echo "[chain $(date '+%H:%M:%S')] DPO done -> eval gate" >> "$CL"
+./stage_eval_50m.sh "$DPO_CKPT" >> "$CL" 2>&1
+echo "[chain $(date '+%H:%M:%S')] eval complete (no merge/release)" >> "$CL"

train/chain_sft_v26_after_evals.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+# Chains SFT v26 after the merge-candidate main-battery evals finish
+# (one heavy torch job at a time). Polls every 60s for all three
+# main-battery .done markers, then launches the SFT v26 watchdog.
+# Launch detached: setsid nohup ./train/chain_sft_v26_after_evals.sh >/dev/null 2>&1 </dev/null & disown
+set -u
+cd "$(dirname "$0")/.."
+LOG=logs/chain_sft_v26.log
+echo "=== chain start $(date '+%Y-%m-%d %H:%M:%S') ===" >> "$LOG"
+
+for i in $(seq 1 600); do
+  done_count=0
+  for c in soup taskarith ties; do
+    if [ -f "logs/eval_cand_merge_${c}_main.done" ]; then
+      done_count=$((done_count+1))
+    fi
+  done
+  if [ "$done_count" -ge 3 ]; then
+    echo "=== all merge evals done at $(date '+%H:%M:%S'); launching SFT v26 ===" >> "$LOG"
+    setsid nohup ./train/watchdog_lora_v26.sh >/dev/null 2>&1 </dev/null & disown
+    echo "SFT v26 watchdog launched" >> "$LOG"
+    exit 0
+  fi
+  sleep 60
+done
+echo "=== chain timeout after 10h; merge evals not all done ===" >> "$LOG"
+exit 1

train/grow_weights.py

@@ -0,0 +1,213 @@
+"""Grow the pretrained 7.8M baseline. Verified path: DEPTH (identity-init new blocks).
+
+Per skill tiny-model-phase2 (measured Aug 2026 on this tablet):
+  - width upscaling 320->512 does NOT transfer (val loss 2.58 -> 6.1-7.7) because
+    RMSNorm/rope/groupnorm/recurrence all depend on d_model.
+  - depth growth with identity blocks DOES preserve baseline exactly (2.567 vs 2.578).
+
+Usage:
+  PYTHONPATH=$PWD .venv/bin/python train/grow_weights.py \
+      --base ckpt/nlp_full --config tiny13m --ckpt ckpt/tiny13m_grown --verify
+"""
+import argparse
+import copy
+import random
+
+import torch
+import torch.nn.functional as F
+
+from model.config import TinyLiquidConfig, CONFIGS
+from model.tiny_liquid import TinyLiquid
+from model.utils import latest_ckpt
+from data.tokenizer import load_tokenizer
+
+STD = 0.02
+NEW_BLOCK_SCALE = 0.1
+
+
+def pad_noise(old, rows_new=None, cols_new=None):
+    old = old.float()
+    r0, c0 = old.shape
+    rn = rows_new if rows_new is not None else r0
+    cn = cols_new if cols_new is not None else c0
+    out = torch.empty(rn, cn)
+    out.fill_(0.0)
+    out[:r0, :c0] = old
+    mask = torch.ones_like(out, dtype=torch.bool)
+    mask[:r0, :c0] = False
+    out[mask] = torch.normal(0.0, STD, size=(int(mask.sum()),))
+    return out
+
+
+def grow_block_width(block, d_new, basis_rows_new, h_new):
+    nb = copy.deepcopy(block)
+    nb["norm1.weight"] = torch.cat([block["norm1.weight"].float(),
+                                    torch.ones(d_new - block["norm1.weight"].shape[0])])
+    nb["basis.w"] = pad_noise(block["basis.w"], rows_new=basis_rows_new, cols_new=d_new)
+    nb["basis.w_forget"] = pad_noise(block["basis.w_forget"], rows_new=basis_rows_new, cols_new=d_new)
+    nb["basis.gn.weight"] = torch.cat([block["basis.gn.weight"].float(),
+                                       torch.ones(basis_rows_new - block["basis.gn.weight"].shape[0])])
+    nb["basis.gn.bias"] = torch.cat([block["basis.gn.bias"].float(),
+                                     torch.zeros(basis_rows_new - block["basis.gn.bias"].shape[0])])
+    nb["norm2.weight"] = torch.cat([block["norm2.weight"].float(),
+                                    torch.ones(d_new - block["norm2.weight"].shape[0])])
+    nb["mlp.up.weight"] = pad_noise(block["mlp.up.weight"], rows_new=h_new, cols_new=d_new)
+    nb["mlp.gate.weight"] = pad_noise(block["mlp.gate.weight"], rows_new=h_new, cols_new=d_new)
+    nb["mlp.forget.weight"] = pad_noise(block["mlp.forget.weight"], rows_new=h_new, cols_new=d_new)
+    nb["mlp.down.weight"] = pad_noise(block["mlp.down.weight"], rows_new=d_new, cols_new=h_new)
+    return nb
+
+
+def identity_block(cfg, prefix):
+    """A block that is an exact identity at init (output == input)."""
+    d = cfg.d_model
+    e = cfg.basis_n * cfg.basis_b
+    h = cfg.mlp_ratio * d
+    sd = {}
+    sd[f"{prefix}.norm1.weight"] = torch.ones(d)
+    sd[f"{prefix}.basis.w"] = torch.zeros(e, d)
+    sd[f"{prefix}.basis.w_forget"] = torch.zeros(e, d)
+    sd[f"{prefix}.basis.gn.weight"] = torch.ones(e)
+    sd[f"{prefix}.basis.gn.bias"] = torch.zeros(e)
+    sd[f"{prefix}.norm2.weight"] = torch.ones(d)
+    sd[f"{prefix}.mlp.up.weight"] = torch.zeros(h, d)
+    sd[f"{prefix}.mlp.gate.weight"] = torch.zeros(h, d)
+    sd[f"{prefix}.mlp.forget.weight"] = torch.zeros(h, d)
+    sd[f"{prefix}.mlp.down.weight"] = torch.zeros(d, h)
+    return sd
+
+
+def grow_depth(sd_old, cfg_old, cfg_new):
+    """Copy trunk exactly; append identity blocks. Baseline loss preserved."""
+    assert cfg_old.d_model == cfg_new.d_model
+    assert cfg_old.basis_n == cfg_new.basis_n and cfg_old.basis_b == cfg_new.basis_b
+    assert cfg_old.mlp_ratio == cfg_new.mlp_ratio
+    assert cfg_new.n_blocks >= cfg_old.n_blocks
+    grown = {}
+    for k, v in sd_old["model"].items():
+        grown[k] = v.clone()
+    for j in range(cfg_old.n_blocks, cfg_new.n_blocks):
+        grown.update(identity_block(cfg_new, f"blocks.{j}"))
+    return grown
+
+
+def grow_width(sd_old, cfg_old, cfg_new):
+    """Width upscaling -- EXPERIMENTAL, does NOT transfer on this architecture."""
+    old = sd_old["model"]
+    d_old, d_new = cfg_old.d_model, cfg_new.d_model
+    basis_old = cfg_old.basis_n * cfg_old.basis_b
+    basis_new = cfg_new.basis_n * cfg_new.basis_b
+    h_old, h_new = cfg_old.mlp_ratio * d_old, cfg_new.mlp_ratio * d_new
+    grown = {}
+    grown["tok_emb.weight"] = pad_noise(old["tok_emb.weight"], cols_new=d_new)
+    grown["persona_emb.weight"] = pad_noise(old["persona_emb.weight"], cols_new=d_new)
+    grown["norm_out.weight"] = torch.cat([old["norm_out.weight"].float(),
+                                          torch.ones(d_new - d_old)])
+    for i in range(cfg_old.n_blocks):
+        block = {k[len(f"blocks.{i}."):]: v for k, v in old.items() if k.startswith(f"blocks.{i}.")}
+        gb = grow_block_width(block, d_new, basis_new, h_new)
+        for k, v in gb.items():
+            grown[f"blocks.{i}.{k}"] = v
+    last = {k[len(f"blocks.{cfg_old.n_blocks-1}."):]: v for k, v in old.items()
+            if k.startswith(f"blocks.{cfg_old.n_blocks-1}.")}
+    seed = grow_block_width(last, d_new, basis_new, h_new)
+    for j in range(cfg_old.n_blocks, cfg_new.n_blocks):
+        for k, v in seed.items():
+            grown[f"blocks.{j}.{k}"] = (v * NEW_BLOCK_SCALE).clone()
+    return grown
+
+
+@torch.no_grad()
+def val_loss(model, tok, data_path, batch=16, seq=256, n_batches=20):
+    import numpy as np
+    tokens = torch.from_numpy(np.fromfile(data_path, dtype="uint16")).long()
+    rng = random.Random(42)
+    total = 0.0
+    for _ in range(n_batches):
+        pos = rng.randrange(0, max(1, len(tokens) - seq - 1))
+        x = tokens[pos:pos + batch * seq].view(batch, seq)
+        y = tokens[pos + 1:pos + 1 + batch * seq].view(batch, seq)
+        logits = model(x)
+        total += F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1)).item()
+    return total / n_batches
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--base", default="ckpt/nlp_full")
+    ap.add_argument("--config", default="tiny13m")
+    ap.add_argument("--mode", choices=["depth", "width", "tower"], default="depth")
+    ap.add_argument("--ckpt", default="ckpt/tiny13m_grown")
+    ap.add_argument("--verify", action="store_true")
+    ap.add_argument("--threads", type=int, default=4)
+    args = ap.parse_args()
+
+    torch.set_num_threads(args.threads)
+    tok = load_tokenizer("data/tokenizer.json")
+    cfg_new = TinyLiquidConfig(vocab_size=tok.get_vocab_size(), **CONFIGS[args.config])
+
+    base_path = latest_ckpt(args.base)
+    assert base_path, f"no checkpoint in {args.base}"
+    sd = torch.load(base_path, map_location="cpu")
+    cfg_old = TinyLiquidConfig(vocab_size=tok.get_vocab_size(),
+                               **{k: v for k, v in sd["config"].items() if k != "vocab_size"})
+    print(f"base: {base_path}  {cfg_old.params_estimate()/1e6:.2f}M -> {cfg_new.params_estimate()/1e6:.2f}M ({args.mode})")
+
+    if args.mode == "depth":
+        grown = grow_depth(sd, cfg_old, cfg_new)
+    elif args.mode == "width":
+        grown = grow_width(sd, cfg_old, cfg_new)
+    else:
+        # tower: keep the trained trunk, add identity-init wide tower.
+        # If tower_d changed, handle by loading trunk + padding tower projection.
+        assert cfg_new.tower_d and cfg_new.tower_blocks, "tower mode needs tower_d/tower_blocks in config"
+        model_old = TinyLiquid(cfg_old)
+        model_old.load_state_dict(sd["model"])
+        model_old.eval()
+        model = TinyLiquid(cfg_new)
+        old_sd = sd["model"]
+        new_sd = dict(model.state_dict())
+        for k, v in old_sd.items():
+            if not k.startswith("up_proj") and not k.startswith("down_proj") and not k.startswith("tower."):
+                if k in new_sd:
+                    new_sd[k] = v.clone()
+                else:
+                    print(f"WARNING: {k} in old but not in new model")
+        if "up_proj" in old_sd:
+            old_up = old_sd["up_proj"]; old_down = old_sd["down_proj"]
+            new_up = new_sd["up_proj"]; new_down = new_sd["down_proj"]
+            if old_up.shape == new_up.shape:
+                new_sd["up_proj"] = old_up.clone()
+                new_sd["down_proj"] = old_down.clone()
+            else:
+                new_sd["up_proj"] = new_sd["up_proj"].clone()
+                new_sd["down_proj"] = new_sd["down_proj"].clone()
+                min_td = min(old_up.shape[0], new_up.shape[0])
+                min_d = min(old_up.shape[1], new_up.shape[1])
+                new_sd["up_proj"][:min_td, :min_d] = old_up[:min_td, :min_d]
+                min_d2 = min(old_down.shape[0], new_down.shape[0])
+                min_td2 = min(old_down.shape[1], new_down.shape[1])
+                new_sd["down_proj"][:min_d2, :min_td2] = old_down[:min_d2, :min_td2]
+        grown = new_sd
+
+    model = TinyLiquid(cfg_new)
+    missing, unexpected = model.load_state_dict(grown, strict=False)
+    assert not missing and not unexpected, (missing, unexpected)
+
+    if args.verify:
+        base_model = TinyLiquid(cfg_old)
+        base_model.load_state_dict(sd["model"])
+        base_model.eval(); model.eval()
+        print(f"baseline val loss (20 bat): {val_loss(base_model, tok, 'data/valid.bin', n_batches=20):.4f}")
+        print(f"grown    val loss (20 bat): {val_loss(model, tok, 'data/valid.bin', n_batches=20):.4f}")
+
+    out_dir = __import__("pathlib").Path(args.ckpt)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    out = out_dir / ("model_final.pt" if args.mode == "tower" else "model_grown.pt")
+    torch.save({"config": cfg_new.__dict__, "model": model.state_dict(),
+                "step": 0, "best_val": float("inf")}, out)
+    print("saved", out)
+
+
+if __name__ == "__main__":
+    main()

train/grow_weights.py.bak

@@ -0,0 +1,190 @@
+"""Grow the pretrained 7.8M baseline. Verified path: DEPTH (identity-init new blocks).
+
+Per skill tiny-model-phase2 (measured Aug 2026 on this tablet):
+  - width upscaling 320->512 does NOT transfer (val loss 2.58 -> 6.1-7.7) because
+    RMSNorm/rope/groupnorm/recurrence all depend on d_model.
+  - depth growth with identity blocks DOES preserve baseline exactly (2.567 vs 2.578).
+
+Usage:
+  PYTHONPATH=$PWD .venv/bin/python train/grow_weights.py \
+      --base ckpt/nlp_full --config tiny13m --ckpt ckpt/tiny13m_grown --verify
+"""
+import argparse
+import copy
+import random
+
+import torch
+import torch.nn.functional as F
+
+from model.config import TinyLiquidConfig, CONFIGS
+from model.tiny_liquid import TinyLiquid
+from model.utils import latest_ckpt
+from data.tokenizer import load_tokenizer
+
+STD = 0.02
+NEW_BLOCK_SCALE = 0.1
+
+
+def pad_noise(old, rows_new=None, cols_new=None):
+    old = old.float()
+    r0, c0 = old.shape
+    rn = rows_new if rows_new is not None else r0
+    cn = cols_new if cols_new is not None else c0
+    out = torch.empty(rn, cn)
+    out.fill_(0.0)
+    out[:r0, :c0] = old
+    mask = torch.ones_like(out, dtype=torch.bool)
+    mask[:r0, :c0] = False
+    out[mask] = torch.normal(0.0, STD, size=(int(mask.sum()),))
+    return out
+
+
+def grow_block_width(block, d_new, basis_rows_new, h_new):
+    nb = copy.deepcopy(block)
+    nb["norm1.weight"] = torch.cat([block["norm1.weight"].float(),
+                                    torch.ones(d_new - block["norm1.weight"].shape[0])])
+    nb["basis.w"] = pad_noise(block["basis.w"], rows_new=basis_rows_new, cols_new=d_new)
+    nb["basis.w_forget"] = pad_noise(block["basis.w_forget"], rows_new=basis_rows_new, cols_new=d_new)
+    nb["basis.gn.weight"] = torch.cat([block["basis.gn.weight"].float(),
+                                       torch.ones(basis_rows_new - block["basis.gn.weight"].shape[0])])
+    nb["basis.gn.bias"] = torch.cat([block["basis.gn.bias"].float(),
+                                     torch.zeros(basis_rows_new - block["basis.gn.bias"].shape[0])])
+    nb["norm2.weight"] = torch.cat([block["norm2.weight"].float(),
+                                    torch.ones(d_new - block["norm2.weight"].shape[0])])
+    nb["mlp.up.weight"] = pad_noise(block["mlp.up.weight"], rows_new=h_new, cols_new=d_new)
+    nb["mlp.gate.weight"] = pad_noise(block["mlp.gate.weight"], rows_new=h_new, cols_new=d_new)
+    nb["mlp.forget.weight"] = pad_noise(block["mlp.forget.weight"], rows_new=h_new, cols_new=d_new)
+    nb["mlp.down.weight"] = pad_noise(block["mlp.down.weight"], rows_new=d_new, cols_new=h_new)
+    return nb
+
+
+def identity_block(cfg, prefix):
+    """A block that is an exact identity at init (output == input)."""
+    d = cfg.d_model
+    e = cfg.basis_n * cfg.basis_b
+    h = cfg.mlp_ratio * d
+    sd = {}
+    sd[f"{prefix}.norm1.weight"] = torch.ones(d)
+    sd[f"{prefix}.basis.w"] = torch.zeros(e, d)
+    sd[f"{prefix}.basis.w_forget"] = torch.zeros(e, d)
+    sd[f"{prefix}.basis.gn.weight"] = torch.ones(e)
+    sd[f"{prefix}.basis.gn.bias"] = torch.zeros(e)
+    sd[f"{prefix}.norm2.weight"] = torch.ones(d)
+    sd[f"{prefix}.mlp.up.weight"] = torch.zeros(h, d)
+    sd[f"{prefix}.mlp.gate.weight"] = torch.zeros(h, d)
+    sd[f"{prefix}.mlp.forget.weight"] = torch.zeros(h, d)
+    sd[f"{prefix}.mlp.down.weight"] = torch.zeros(d, h)
+    return sd
+
+
+def grow_depth(sd_old, cfg_old, cfg_new):
+    """Copy trunk exactly; append identity blocks. Baseline loss preserved."""
+    assert cfg_old.d_model == cfg_new.d_model
+    assert cfg_old.basis_n == cfg_new.basis_n and cfg_old.basis_b == cfg_new.basis_b
+    assert cfg_old.mlp_ratio == cfg_new.mlp_ratio
+    assert cfg_new.n_blocks >= cfg_old.n_blocks
+    grown = {}
+    for k, v in sd_old["model"].items():
+        grown[k] = v.clone()
+    for j in range(cfg_old.n_blocks, cfg_new.n_blocks):
+        grown.update(identity_block(cfg_new, f"blocks.{j}"))
+    return grown
+
+
+def grow_width(sd_old, cfg_old, cfg_new):
+    """Width upscaling -- EXPERIMENTAL, does NOT transfer on this architecture."""
+    old = sd_old["model"]
+    d_old, d_new = cfg_old.d_model, cfg_new.d_model
+    basis_old = cfg_old.basis_n * cfg_old.basis_b
+    basis_new = cfg_new.basis_n * cfg_new.basis_b
+    h_old, h_new = cfg_old.mlp_ratio * d_old, cfg_new.mlp_ratio * d_new
+    grown = {}
+    grown["tok_emb.weight"] = pad_noise(old["tok_emb.weight"], cols_new=d_new)
+    grown["persona_emb.weight"] = pad_noise(old["persona_emb.weight"], cols_new=d_new)
+    grown["norm_out.weight"] = torch.cat([old["norm_out.weight"].float(),
+                                          torch.ones(d_new - d_old)])
+    for i in range(cfg_old.n_blocks):
+        block = {k[len(f"blocks.{i}."):]: v for k, v in old.items() if k.startswith(f"blocks.{i}.")}
+        gb = grow_block_width(block, d_new, basis_new, h_new)
+        for k, v in gb.items():
+            grown[f"blocks.{i}.{k}"] = v
+    last = {k[len(f"blocks.{cfg_old.n_blocks-1}."):]: v for k, v in old.items()
+            if k.startswith(f"blocks.{cfg_old.n_blocks-1}.")}
+    seed = grow_block_width(last, d_new, basis_new, h_new)
+    for j in range(cfg_old.n_blocks, cfg_new.n_blocks):
+        for k, v in seed.items():
+            grown[f"blocks.{j}.{k}"] = (v * NEW_BLOCK_SCALE).clone()
+    return grown
+
+
+@torch.no_grad()
+def val_loss(model, tok, data_path, batch=16, seq=256, n_batches=20):
+    import numpy as np
+    tokens = torch.from_numpy(np.fromfile(data_path, dtype="uint16")).long()
+    rng = random.Random(42)
+    total = 0.0
+    for _ in range(n_batches):
+        pos = rng.randrange(0, max(1, len(tokens) - seq - 1))
+        x = tokens[pos:pos + batch * seq].view(batch, seq)
+        y = tokens[pos + 1:pos + 1 + batch * seq].view(batch, seq)
+        logits = model(x)
+        total += F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1)).item()
+    return total / n_batches
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--base", default="ckpt/nlp_full")
+    ap.add_argument("--config", default="tiny13m")
+    ap.add_argument("--mode", choices=["depth", "width", "tower"], default="depth")
+    ap.add_argument("--ckpt", default="ckpt/tiny13m_grown")
+    ap.add_argument("--verify", action="store_true")
+    ap.add_argument("--threads", type=int, default=4)
+    args = ap.parse_args()
+
+    torch.set_num_threads(args.threads)
+    tok = load_tokenizer("data/tokenizer.json")
+    cfg_new = TinyLiquidConfig(vocab_size=tok.get_vocab_size(), **CONFIGS[args.config])
+
+    base_path = latest_ckpt(args.base)
+    assert base_path, f"no checkpoint in {args.base}"
+    sd = torch.load(base_path, map_location="cpu")
+    cfg_old = TinyLiquidConfig(vocab_size=tok.get_vocab_size(),
+                               **{k: v for k, v in sd["config"].items() if k != "vocab_size"})
+    print(f"base: {base_path}  {cfg_old.params_estimate()/1e6:.2f}M -> {cfg_new.params_estimate()/1e6:.2f}M ({args.mode})")
+
+    if args.mode == "depth":
+        grown = grow_depth(sd, cfg_old, cfg_new)
+    elif args.mode == "width":
+        grown = grow_width(sd, cfg_old, cfg_new)
+    else:
+        # tower: keep the trained trunk, add identity-init wide tower (up=identity,
+        # tower blocks identity, down=zero) => baseline output preserved exactly.
+        assert cfg_new.tower_d and cfg_new.tower_blocks, "tower mode needs tower_d/tower_blocks in config"
+        model = TinyLiquid(cfg_new)
+        miss, unexp = model.load_state_dict(sd["model"], strict=False)
+        assert not unexp, unexp
+        assert all("tower" in k or "up_proj" in k or "down_proj" in k for k in miss), miss
+        grown = model.state_dict()
+
+    model = TinyLiquid(cfg_new)
+    missing, unexpected = model.load_state_dict(grown, strict=False)
+    assert not missing and not unexpected, (missing, unexpected)
+
+    if args.verify:
+        base_model = TinyLiquid(cfg_old)
+        base_model.load_state_dict(sd["model"])
+        base_model.eval(); model.eval()
+        print(f"baseline val loss (20 bat): {val_loss(base_model, tok, 'data/valid.bin', n_batches=20):.4f}")
+        print(f"grown    val loss (20 bat): {val_loss(model, tok, 'data/valid.bin', n_batches=20):.4f}")
+
+    out_dir = __import__("pathlib").Path(args.ckpt)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    out = out_dir / ("model_final.pt" if args.mode == "tower" else "model_grown.pt")
+    torch.save({"config": cfg_new.__dict__, "model": model.state_dict(),
+                "step": 0, "best_val": float("inf")}, out)
+    print("saved", out)
+
+
+if __name__ == "__main__":
+    main()

train/map_vocab.py

@@ -0,0 +1,177 @@
+"""Expand tokenizer vocab on a trained checkpoint (8k -> 16k), baseline-preserving.
+
+The lm_head is TIED (TinyLiquid cfg.tie_embeddings=True), so tok_emb.weight
+is the ONLY vocab-sized tensor. We map old token ids -> new token ids by
+DECODED-TEXT EXACT MATCH: old rows whose text exists as a single new token are
+copied exactly (their behavior is preserved); others fall back to the first
+token of the re-encoded text or deterministic noise (std 0.02).
+
+Usage:
+  .venv/bin/python train/map_vocab.py \
+      --base ckpt/hybrid50m_pretrain --old-tok data/tokenizer.json \
+      --new-tok data/tokenizer16k.json --config hybrid50m \
+      --out ckpt/hybrid50m_v16k_init.pt
+"""
+import argparse
+import random
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+from model.config import TinyLiquidConfig
+from model.tiny_liquid import TinyLiquid
+from model.utils import latest_ckpt
+from data.tokenizer import load_tokenizer, SPECIAL_TOKENS
+
+STD = 0.02
+
+
+def build_mapping(old_tok, new_tok):
+    """Return (old_id->new_id|None dict, set of exact-match old ids).
+
+    Special tokens map by index (identical set, identical order, identical
+    ids in the BPE trainer). Regular tokens map by decoded-text exact match,
+    else by the first new token of the re-encoded text.
+    """
+    old_vocab, new_vocab = old_tok.get_vocab_size(), new_tok.get_vocab_size()
+    new_text_ids = {}
+    for j in range(new_vocab):
+        t = new_tok.decode([j])
+        if t and t not in new_text_ids:
+            new_text_ids[t] = j
+
+    mapping, exact_ids = {}, set()
+    for i in range(old_vocab):
+        if i < len(SPECIAL_TOKENS):
+            mapping[i] = i
+            exact_ids.add(i)
+            continue
+        text = old_tok.decode([i])
+        if not text:
+            mapping[i] = None
+            continue
+        exact = new_text_ids.get(text)
+        if exact is not None:
+            mapping[i] = exact
+            exact_ids.add(i)
+            continue
+        ids = new_tok.encode(text).ids
+        first = [j for j in ids if j >= len(SPECIAL_TOKENS)]
+        mapping[i] = first[0] if first else None
+    return mapping, exact_ids
+
+
+def expand_checkpoint(sd, cfg_new, mapping, exact_ids, seed=42):
+    """Copy every weight; expand tok_emb rows (exact copy -> fallback -> noise)."""
+    new_model = TinyLiquid(cfg_new)
+    new_sd = dict(new_model.state_dict())
+    for k, v in sd["model"].items():
+        if k == "tok_emb.weight":
+            continue
+        assert k in new_sd, f"key {k} missing in new model"
+        new_sd[k] = v.clone()
+
+    g = torch.Generator().manual_seed(seed)
+    old_emb = sd["model"]["tok_emb.weight"].float()
+    emb = torch.empty(cfg_new.vocab_size, old_emb.shape[1])
+    emb.normal_(0.0, STD, generator=g)
+    exact = partial = noise = 0
+    for i in range(old_emb.shape[0]):
+        tgt = mapping.get(i)
+        if tgt is None:
+            noise += 1
+            continue
+        emb[tgt] = old_emb[i]
+        if i in exact_ids:
+            exact += 1
+        else:
+            partial += 1
+    new_sd["tok_emb.weight"] = emb
+
+    model = TinyLiquid(cfg_new)
+    missing, unexpected = model.load_state_dict(new_sd, strict=True)
+    assert not missing and not unexpected, (missing, unexpected)
+    return new_sd, {"exact": exact, "partial": partial, "noise": noise}
+
+
+@torch.no_grad()
+def val_loss(model, data_path, batch=8, seq=256, n_batches=20):
+    import numpy as np
+    tokens = torch.from_numpy(np.fromfile(data_path, dtype="uint16")).long()
+    rng = random.Random(42)
+    total = 0.0
+    for _ in range(n_batches):
+        pos = rng.randrange(0, max(1, len(tokens) - seq - 1))
+        x = tokens[pos:pos + batch * seq].view(batch, seq)
+        y = tokens[pos + 1:pos + 1 + batch * seq].view(batch, seq)
+        logits = model(x)
+        total += F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1)).item()
+    return total / n_batches
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--base", default="ckpt/hybrid50m_pretrain")
+    ap.add_argument("--old-tok", default="data/tokenizer.json")
+    ap.add_argument("--new-tok", default="data/tokenizer16k.json")
+    ap.add_argument("--config", default="hybrid50m")
+    ap.add_argument("--out", default="ckpt/hybrid50m_v16k_init.pt")
+    ap.add_argument("--old-val", default="data/valid.bin")
+    ap.add_argument("--new-val", default="data/valid16k.bin")
+    ap.add_argument("--threads", type=int, default=4)
+    ap.add_argument("--verify", action="store_true",
+                    help="measure val-loss delta vs baseline (slower)")
+    args = ap.parse_args()
+
+    torch.set_num_threads(args.threads)
+    old_tok = load_tokenizer(args.old_tok)
+    new_tok = load_tokenizer(args.new_tok)
+    old_vocab = old_tok.get_vocab_size()
+    new_vocab = new_tok.get_vocab_size()
+    assert new_vocab > old_vocab, f"new vocab {new_vocab} must exceed old {old_vocab}"
+
+    base_path = Path(args.base)
+    path = base_path if base_path.is_file() else latest_ckpt(args.base)
+    assert path, f"no checkpoint at {args.base}"
+    sd = torch.load(str(path), map_location="cpu")
+    cfg_old = TinyLiquidConfig(**sd["config"])
+    cfg_new = TinyLiquidConfig(vocab_size=new_vocab,
+                               **{k: v for k, v in sd["config"].items()
+                                  if k != "vocab_size"})
+    print(f"base: {path}  vocab {cfg_old.vocab_size} -> {cfg_new.vocab_size}, "
+          f"{cfg_old.params_estimate()/1e6:.2f}M -> {cfg_new.params_estimate()/1e6:.2f}M")
+
+    mapping, exact_ids = build_mapping(old_tok, new_tok)
+    new_sd, stats = expand_checkpoint(sd, cfg_new, mapping, exact_ids)
+    total = old_vocab
+    print(f"token map: exact={stats['exact']} partial={stats['partial']} "
+          f"noise={stats['noise']} ({100*stats['exact']/total:.1f}% exact)")
+
+    if args.verify:
+        model_old = TinyLiquid(cfg_old)
+        model_old.load_state_dict(sd["model"])
+        model_old.eval()
+        model = TinyLiquid(cfg_new)
+        model.load_state_dict(new_sd)
+        model.eval()
+        lo = val_loss(model_old, args.old_val, n_batches=20)
+        ln = val_loss(model, args.new_val, n_batches=20)
+        print(f"old-tokenizer baseline loss {lo:.4f} -> new-tokenizer 16k loss {ln:.4f} "
+              f"(delta {ln-lo:+.4f})", flush=True)
+        # A tokenizer change changes the target sequence. Exact loss parity
+        # before new merges are trained is not a valid gate; require finite
+        # loss and strong coverage, then gate recovery during continuation.
+        assert torch.isfinite(torch.tensor(lo)) and torch.isfinite(torch.tensor(ln)), "non-finite mapping loss"
+        assert stats["exact"] / total >= 0.60, "exact token coverage gate failed"
+
+    out = Path(args.out)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    torch.save({"config": cfg_new.__dict__, "model": new_sd, "step": 0,
+                "best_val": sd.get("best_val", float("inf")), "tag": "vocab16k-init"},
+               out)
+    print("saved", out)
+
+
+if __name__ == "__main__":
+    main()

train/parallel_merges.py

@@ -0,0 +1,101 @@
+"""Parallel merge recipes for post-training checkpoints (tiny-model-posttrain).
+
+Runs model soup (arXiv 2203.05482) and task arithmetic (arXiv 2212.04089)
+on the SAME base, producing one checkpoint per recipe. TIES is handled by
+train/ties_merge.py (arXiv 2306.01708). After merging, battery-eval every
+candidate and keep the best (LFM2 2511.23404 §4.4: parallel apply -> eval ->
+select). Never naive-average adapters; these recipes operate on folded
+full-weight checkpoints where delta = task_ckpt - base is the true task
+vector.
+
+Usage:
+  .venv/bin/python train/parallel_merges.py \
+      --base ckpt/hybrid50m_v16k_pretrain/model_5000.pt \
+      --tasks ckpt/hybrid50m_v25_lora/best.pt ckpt/hybrid50m_v25_dpo/model_final.pt \
+      --out-dir ckpt/hybrid50m_v25_merges \
+      --lambda-ta 0.5
+"""
+import argparse
+from pathlib import Path
+
+import torch
+
+from model.config import TinyLiquidConfig
+from model.tiny_liquid import TinyLiquid
+from model.utils import latest_ckpt
+
+
+def load_sd(path, tag):
+    sd = torch.load(str(path), map_location="cpu", weights_only=False)
+    print(f"  {tag}: {path}  step={sd.get('step', '?')}  tag={sd.get('tag', '-')}", flush=True)
+    return sd
+
+
+def model_soup(task_sds):
+    """Simple average of task weights (all tasks trained from the same base)."""
+    soup = {}
+    keys = [k for k in task_sds[0] if task_sds[0][k].is_floating_point()
+            and all(k in sd for sd in task_sds[1:])]
+    for k in keys:
+        soup[k] = torch.stack([sd[k].float() for sd in task_sds]).mean(dim=0)
+    return soup
+
+
+def task_arithmetic(base_sd, task_sds, lam):
+    """base + lam * sum(task_i - base)."""
+    merged = {}
+    keys = [k for k in base_sd if base_sd[k].is_floating_point()
+            and all(k in sd for sd in task_sds)]
+    for k in keys:
+        base = base_sd[k].float()
+        delta = torch.zeros_like(base)
+        for sd in task_sds:
+            delta += sd[k].float() - base
+        merged[k] = base + lam * delta
+    return merged
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--base", required=True)
+    ap.add_argument("--tasks", nargs="+", required=True)
+    ap.add_argument("--out-dir", required=True)
+    ap.add_argument("--lambda-ta", type=float, default=0.5)
+    ap.add_argument("--threads", type=int, default=4)
+    args = ap.parse_args()
+    assert len(args.tasks) >= 2, "merges need >= 2 task checkpoints"
+
+    torch.set_num_threads(args.threads)
+    base_path = Path(args.base)
+    base_ckpt = base_path if base_path.is_file() else latest_ckpt(args.base)
+    base = load_sd(base_ckpt, "base")
+    base_sd = base["model"]
+
+    task_sds = []
+    for i, t in enumerate(args.tasks):
+        tp = Path(t)
+        tp = tp if tp.is_file() else latest_ckpt(t)
+        task_sds.append(load_sd(tp, f"task{i}")["model"])
+
+    recipes = {
+        "soup": model_soup(task_sds),
+        f"taskarith_l{args.lambda_ta}".replace(".", "p"): task_arithmetic(base_sd, task_sds, args.lambda_ta),
+    }
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    for name, merged in recipes.items():
+        cfg = TinyLiquidConfig(**base["config"])
+        cfg.mtp_heads = 0
+        model = TinyLiquid(cfg)
+        missing, unexpected = model.load_state_dict(merged, strict=False)
+        if missing or unexpected:
+            print(f"[{name}] ignored {len(missing)} missing / {len(unexpected)} unexpected keys", flush=True)
+        out = out_dir / f"{name}.pt"
+        torch.save({"config": base["config"], "model": merged, "step": 0,
+                    "best_val": base.get("best_val", float("inf")),
+                    "tag": f"{name}-{len(task_sds)}tasks"}, out)
+        print("saved", out, flush=True)
+
+
+if __name__ == "__main__":
+    main()

train/ties_merge.py

@@ -0,0 +1,105 @@
+"""TIES task-vector merge for post-training checkpoints (tiny-model-posttrain).
+
+After SFT/DPO, merge the task vectors (ckpt_i - base) instead of naive
+averaging (measured failure at 25M):
+  1. TRIM: keep the top-`keep` fraction of each tensor by |delta|.
+  2. SIGN CONSENSUS: where task vectors disagree on sign, zero the delta.
+  3. MERGE: mean of the surviving deltas, added back to the base.
+
+Pure function `ties_merge` is unit-tested; the CLI wraps it for the repo's
+folded state-dict checkpoints ({"model": ..., "config": ...}).
+
+Usage:
+  .venv/bin/python train/ties_merge.py \
+      --base ckpt/hybrid50m_pretrain \
+      --tasks ckpt/hybrid50m_sft_a ckpt/hybrid50m_dpo_a ckpt/hybrid50m_dpo_b \
+      --keep 0.2 --out ckpt/hybrid50m_ties.pt
+"""
+import argparse
+import copy
+from pathlib import Path
+
+import torch
+
+from model.config import TinyLiquidConfig
+from model.tiny_liquid import TinyLiquid
+from model.utils import latest_ckpt
+
+
+def trim_delta(delta, keep=0.2):
+    """Zero all but the top-`keep` fraction of |delta| (per tensor)."""
+    d = delta.clone().float()
+    flat = d.abs().flatten()
+    k = max(1, int(round(flat.numel() * keep)))
+    thresh = flat.topk(k).values.min()
+    d[d.abs() < thresh] = 0.0
+    return d
+
+
+def ties_merge(base_sd, task_sds, keep=0.2):
+    """Return merged state dict: base + consensus mean of trimmed deltas."""
+    merged = {}
+    keys = [k for k in base_sd if base_sd[k].is_floating_point()
+            and all(k in t for t in task_sds)]
+    for k in keys:
+        base = base_sd[k].float()
+        deltas = [task[k].float() - base for task in task_sds]
+        trimmed = [trim_delta(d, keep) for d in deltas]
+        signs = torch.stack([t.sign() for t in trimmed])
+        consensus = signs.mean(dim=0)
+        agree = consensus.abs() >= 1.0          # all signs agree
+        final = torch.zeros_like(base)
+        if agree.any():
+            surv = torch.stack([t for t in trimmed])
+            final[agree] = surv[:, agree].mean(dim=0)
+        merged[k] = base + final
+    return merged
+
+
+def load_sd(path, tag):
+    sd = torch.load(str(path), map_location="cpu", weights_only=False)
+    print(f"  {tag}: {path}  step={sd.get('step', '?')}  tag={sd.get('tag', '-')}")
+    return sd["model"]
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--base", required=True)
+    ap.add_argument("--tasks", nargs="+", required=True)
+    ap.add_argument("--keep", type=float, default=0.2)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--threads", type=int, default=4)
+    args = ap.parse_args()
+    assert len(args.tasks) >= 2, "TIES needs >= 2 task checkpoints"
+
+    torch.set_num_threads(args.threads)
+    base_path = Path(args.base)
+    base_ckpt = base_path if base_path.is_file() else latest_ckpt(args.base)
+    base = torch.load(str(base_ckpt), map_location="cpu", weights_only=False)
+    base_sd = base["model"]
+    print(f"base: {base_ckpt}")
+
+    task_sds = []
+    for i, t in enumerate(args.tasks):
+        tp = Path(t)
+        tp = tp if tp.is_file() else latest_ckpt(t)
+        task_sds.append(load_sd(tp, f"task{i}"))
+
+    merged = ties_merge(base_sd, task_sds, keep=args.keep)
+    cfg = TinyLiquidConfig(**base["config"])
+    cfg.mtp_heads = 0  # final merged model carries no MTP (pretrain-only)
+    model = TinyLiquid(cfg)
+    missing, unexpected = model.load_state_dict(merged, strict=False)
+    if missing or unexpected:
+        print(f"[ties] ignored {len(missing)} missing / {len(unexpected)} unexpected keys (expected from MTP drift)")
+
+    out = Path(args.out)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    torch.save({"config": base["config"], "model": merged, "step": 0,
+                "best_val": base.get("best_val", float("inf")),
+                "tag": f"ties-{len(task_sds)}tasks-k{args.keep}"}, out)
+    print("saved", out)
+
+
+if __name__ == "__main__":
+    main()

train/train_classifier.py

@@ -0,0 +1,182 @@
+"""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()

train/train_dpo.py

@@ -0,0 +1,288 @@
+"""DPO (direct preference optimization) for TinyLiquid persona naturalness.
+
+Teacher (assistant) writes preference pairs: same prompt, one natural
+persona response (chosen) and one stiff/template response (rejected).
+The model learns to prefer the chosen style at the token level.
+
+Loss: -log sigmoid( beta * (reward_w - reward_r) ), where
+reward = logpi(y|x) - logpi_ref(y|x), averaged over assistant tokens.
+
+Batched forward passes, periodic checkpoints, and --resume so a long
+preference run survives a tablet death. Recipe from tiny-model-preference:
+
+  PYTHONPATH=$PWD .venv/bin/python train/train_dpo.py --base <best_sft> \
+      --data data/prefs_p_all.jsonl --ckpt ckpt/tiny25m_dpo3 \
+      --epochs 1 --batch 4 --seq 512 --lr 4e-6 --beta 0.05 --threads 8
+"""
+
+import argparse
+import json
+import math
+import random
+import sys
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+from model.config import TinyLiquidConfig, CONFIGS
+from model.tiny_liquid import TinyLiquid
+from model.utils import latest_ckpt
+from data.tokenizer import load_tokenizer
+from train.train_lora import wrap_lora, fold_state_dict, val_ppl
+
+PERSONA_T = {"analyst": "<|analyst|>", "skeptic": "<|skeptic|>"}
+
+
+def parse_args():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--base", default="ckpt/distill")
+    ap.add_argument("--data", default="data/prefs_persona.jsonl")
+    ap.add_argument("--tok", default="data/tokenizer.json")
+    ap.add_argument("--ckpt", default="ckpt/dpo")
+    ap.add_argument("--epochs", type=int, default=5)
+    ap.add_argument("--batch", type=int, default=4)
+    ap.add_argument("--seq", type=int, default=256)
+    ap.add_argument("--lr", type=float, default=1e-5)
+    ap.add_argument("--beta", type=float, default=0.1)
+    ap.add_argument("--save-every", type=int, default=100)
+    ap.add_argument("--resume", default=None, help="continue from this checkpoint (model + opt + step); reference policy stays --base")
+    ap.add_argument("--resume-best-ppl", type=float, default=None, help="required fallback best PPL for legacy checkpoints without metadata")
+    ap.add_argument("--reset-optimizer", action="store_true", help="resume model/step but initialize a fresh optimizer")
+    ap.add_argument("--seed", type=int, default=9)
+    ap.add_argument("--threads", type=int, default=8)
+    ap.add_argument("--lora", action="store_true", help="adapter-only DPO: freeze base, train LoRA adapters")
+    ap.add_argument("--r", type=int, default=16)
+    ap.add_argument("--alpha", type=float, default=32.0)
+    ap.add_argument("--dropout", type=float, default=0.05)
+    ap.add_argument("--ppl-guard", type=float, default=60.0)
+    ap.add_argument("--val-bin", default="data/valid.bin")
+    ap.add_argument("--eval-every", type=int, default=100)
+    ap.add_argument("--dpo-weight", type=float, default=1.0,
+                    help="LFM2 joint objective: weight w on the DPO term (default 1.0)")
+    ap.add_argument("--margin", type=float, default=0.1,
+                    help="LFM2 joint objective: margin m inside f(Delta - m) (default 0.1; 0 = plain DPO)")
+    ap.add_argument("--apo-weight", type=float, default=0.2,
+                    help="LFM2 joint objective: weight lambda on the APO-zero absolute term (default 0.2; 0 = plain DPO)")
+    ap.add_argument("--lr-schedule", choices=["constant", "cosine"], default="constant",
+                    help="LR schedule (LFM2 Table 5 uses cosine 8e-7 -> 8e-8 with 0.01 warmup)")
+    ap.add_argument("--lr-min", type=float, default=8e-8,
+                    help="cosine schedule minimum LR (LFM2 Table 5 default 8e-8)")
+    ap.add_argument("--warmup-frac", type=float, default=0.01,
+                    help="fraction of total steps used for linear warmup (LFM2 Table 5 default 0.01)")
+    return ap.parse_args()
+
+
+def lr_at(step, total, lr_max, lr_min, warmup_frac, schedule):
+    """LFM2 Table 5 schedule: linear warmup then cosine decay.
+
+    Constant schedule keeps the historical fixed-LR behavior (default).
+    """
+    if schedule == "constant":
+        return lr_max
+    warmup_steps = max(1, int(round(total * warmup_frac)))
+    if step <= warmup_steps:
+        return lr_min + (lr_max - lr_min) * (step / warmup_steps)
+    progress = (step - warmup_steps) / max(1, total - warmup_steps)
+    return lr_min + 0.5 * (lr_max - lr_min) * (1.0 + math.cos(math.pi * min(1.0, progress)))
+
+
+def build_tensors(tok, persona, prompt, response, seq):
+    p_token = PERSONA_T.get(persona, PERSONA_T["analyst"])
+    u_id, a_id, eot = (tok.token_to_id(t) for t in ("<|user|>", "<|assistant|>", "<|endoftext|>"))
+    parts = [tok.encode(p_token).ids, [u_id], tok.encode(prompt).ids, [a_id],
+             tok.encode(response).ids, [eot]]
+    ids = [i for part in parts for i in part]
+    if len(ids) > seq:
+        ids = ids[: seq - 1] + [eot]
+    asst_start = len(parts[0]) + 1 + len(parts[2]) + 1
+    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
+    return x, y[: x.shape[0]], mask[: x.shape[0]]
+
+
+def build_batch(tok, examples, key, seq):
+    built = [build_tensors(tok, ex.get("persona", "analyst"), ex["prompt"], ex[key], seq) for ex in examples]
+    T = max(x.shape[0] for x, _, _ in built)
+    B = len(built)
+    xs = torch.zeros(B, T, dtype=torch.long)
+    ys = torch.zeros(B, T, dtype=torch.long)
+    ms = torch.zeros(B, T, dtype=torch.bool)
+    for i, (x, y, m) in enumerate(built):
+        xs[i, : x.shape[0]] = x
+        ys[i, : y.shape[0]] = y
+        ms[i, : m.shape[0]] = m
+    return xs, ys, ms
+
+
+def batch_seq_logprob(model, xs, ys, masks, persona_ids):
+    logits = model(xs, persona_ids=persona_ids)
+    logp = F.log_softmax(logits.float(), dim=-1)
+    per_tok = logp.gather(-1, ys.unsqueeze(-1)).squeeze(-1)
+    return (per_tok * masks).sum(dim=-1) / masks.sum(dim=-1).clamp(min=1.0)
+
+
+def lfm2_loss(reward_w, reward_r, beta, margin=0.1, dpo_weight=1.0, apo_weight=0.2):
+    """LFM2 length-normalized joint objective (arXiv 2511.23404, section 4.3).
+
+    L = -E[ w*f(Delta - m) + lambda*g(delta) ]
+      Delta = beta*rew_w - beta*rew_r            (length-normalized relative)
+      delta = sigmoid(beta*rew_w) - sigmoid(beta*rew_r)   (absolute)
+      f = log_sigmoid (DPO),  g = identity (APO-zero)
+    The minus sits outside the expectation, so the APO contribution is
+    -lambda*(sigmoid(scaled_w) - sigmoid(scaled_r)) and can be negative
+    when the chosen response is favored (gradient still widens the gap).
+    Plain DPO is the special case margin=0, apo_weight=0.
+    """
+    scaled_w = beta * reward_w
+    scaled_r = beta * reward_r
+    dpo_term = -F.logsigmoid(scaled_w - scaled_r - margin)
+    apo_term = torch.sigmoid(scaled_r) - torch.sigmoid(scaled_w)
+    return (dpo_weight * dpo_term + apo_weight * apo_term).mean()
+
+
+def main():
+    args = parse_args()
+    torch.set_num_threads(args.threads)
+    torch.manual_seed(args.seed)
+    random.seed(args.seed)
+
+    tok = load_tokenizer(args.tok)
+    base_path = Path(args.base)
+    base_ckpt = base_path if base_path.is_file() else latest_ckpt(args.base)
+    sd = torch.load(base_ckpt, map_location="cpu")
+    config = sd.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"})
+    cfg.mtp_heads = 0  # MTP is pretrain-only; post-training has no MTP heads
+    model = TinyLiquid(cfg)
+    resume = torch.load(args.resume, map_location="cpu") if args.resume else None
+    ref = TinyLiquid(cfg)
+    ref.load_state_dict(sd["model"], strict=False)
+    for p in ref.parameters():
+        p.requires_grad_(False)
+    ref.eval()
+    wrapped = []
+    if resume is not None:
+        if args.lora:
+            wrapped = wrap_lora(model, args.r, args.alpha, args.dropout)
+        model.load_state_dict(resume["model"], strict=False)
+    else:
+        model.load_state_dict(sd["model"], strict=False)
+        if args.lora:
+            wrapped = wrap_lora(model, args.r, args.alpha, args.dropout)
+    if args.lora:
+        for p in model.parameters():
+            p.requires_grad = False
+        for p in model.persona_emb.parameters():
+            p.requires_grad = True
+        for _, lora in wrapped:
+            lora.lora_a.requires_grad = True
+            lora.lora_b.requires_grad = True
+        print(f"adapter-only DPO: {len(wrapped)} lora adapters, "
+              f"trainable {sum(p.numel() for p in model.parameters() if p.requires_grad):,}", flush=True)
+    model.train()
+    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad],
+                            lr=args.lr, betas=(0.9, 0.95), weight_decay=0.02)
+    start_step = 0
+    if resume is not None and not args.reset_optimizer:
+        opt.load_state_dict(resume["opt"])
+        start_step = resume.get("step", 0)
+    if resume is not None:
+        if "best_ppl" in resume and resume["best_ppl"] is not None: best_ppl = float(resume["best_ppl"])
+        elif args.resume_best_ppl is not None: best_ppl = args.resume_best_ppl
+        else: raise ValueError("resume checkpoint lacks best_ppl; pass --resume-best-ppl from the recorded evaluation")
+        if "torch_rng" in resume: torch.set_rng_state(resume["torch_rng"])
+        if "python_rng" in resume: random.setstate(resume["python_rng"])
+    else: best_ppl = float("inf")
+
+    examples = [json.loads(l) for l in open(args.data, encoding="utf-8") if l.strip()]
+    assert examples, "no preference data"
+    print(f"loaded {len(examples)} preference pairs; base {base_ckpt}", flush=True)
+
+    out_dir = Path(args.ckpt)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    steps_per_epoch = max(1, len(examples) // args.batch)
+    total = steps_per_epoch * args.epochs
+    step = 0
+    if start_step:
+        print(f"resuming from {args.resume} at step {start_step}/{total}", flush=True)
+    for ep in range(args.epochs):
+        # Rebuild the same epoch order after resume; global RNG state alone is
+        # insufficient because shuffle happens before skipped steps.
+        epoch_rng = random.Random(args.seed + ep)
+        epoch_rng.shuffle(examples)
+        usable = len(examples) - len(examples) % args.batch
+        for i in range(0, usable, args.batch):
+            step += 1
+            if step <= start_step:
+                continue
+            batch = examples[i : i + args.batch]
+            persona_ids = torch.tensor([2 if ex.get("persona") == "skeptic" else 1 for ex in batch])
+            xw, yw, mw = build_batch(tok, batch, "chosen", args.seq)
+            xr, yr, mr = build_batch(tok, batch, "rejected", args.seq)
+            opt.zero_grad(set_to_none=True)
+            lw = batch_seq_logprob(model, xw, yw, mw, persona_ids)
+            lr = batch_seq_logprob(model, xr, yr, mr, persona_ids)
+            with torch.no_grad():
+                rw = batch_seq_logprob(ref, xw, yw, mw, persona_ids)
+                rr = batch_seq_logprob(ref, xr, yr, mr, persona_ids)
+            reward_w = lw - rw
+            reward_r = lr - rr
+            loss = lfm2_loss(reward_w, reward_r, args.beta,
+                             margin=args.margin,
+                             dpo_weight=args.dpo_weight,
+                             apo_weight=args.apo_weight)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            opt.param_groups[0]["lr"] = lr_at(step, total, args.lr, args.lr_min,
+                                              args.warmup_frac, args.lr_schedule)
+            opt.step()
+            if step % 10 == 0 or step == total:
+                print(f"step {step}/{total} dpo_loss {loss.item():.4f}", flush=True)
+            def save_ck(tag):
+                if wrapped:
+                    torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": step,
+                                 "config": cfg.__dict__, "best_ppl": best_ppl,
+                                 "torch_rng": torch.get_rng_state(), "python_rng": random.getstate()}, str(out_dir / "raw_latest.pt"))
+                sd_save = fold_state_dict(model.state_dict(), wrapped) if wrapped else model.state_dict()
+                torch.save({"model": sd_save, "opt": opt.state_dict(), "step": step,
+                                 "config": cfg.__dict__, "best_ppl": best_ppl,
+                                 "torch_rng": torch.get_rng_state(), "python_rng": random.getstate()}, str(out_dir / tag))
+            if step % args.eval_every == 0:
+                ppl = val_ppl(model, args.val_bin, batch=2, seq=64, n_batches=1, seed=args.seed + step)
+                mark = ""
+                if ppl > args.ppl_guard:
+                    failure = out_dir / "dpo_failed.json"
+                    temp = failure.with_suffix(".partial")
+                    temp.write_text(json.dumps({"status": "ppl-guard-breach", "step": step,
+                                                  "val_ppl": ppl, "ppl_guard": args.ppl_guard}, indent=2) + "\n")
+                    temp.replace(failure)
+                    print(f"  [eval {step}] val_ppl {ppl:.2f} guard {args.ppl_guard} [BREACH]", flush=True)
+                    print("PPL GUARD BREACH — aborting before any checkpoint is marked best", flush=True)
+                    sys.exit(3)
+                if ppl < best_ppl:
+                    best_ppl = ppl
+                    save_ck("best_ppl.pt")
+                    mark = " [best ppl]"
+                if ppl <= args.ppl_guard:
+                    save_ck(f"model_{step}.pt")
+                print(f"  [eval {step}] val_ppl {ppl:.2f} guard {args.ppl_guard}{mark}", flush=True)
+            if step % args.save_every == 0:
+                save_ck(f"model_{step}.pt")
+                print(f"  saved model_{step}.pt", flush=True)
+    save_ck("model_final.pt")
+    completion = out_dir / "dpo_complete.json"
+    temp = completion.with_suffix(".partial")
+    temp.write_text(json.dumps({"status": "complete", "step": step,
+                                "best_ppl": best_ppl,
+                                "checkpoint": "model_final.pt"}, indent=2) + "\n")
+    temp.replace(completion)
+    (out_dir / "dpo_failed.json").unlink(missing_ok=True)
+    print(f"done -> {out_dir}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

train/train_judge.py

@@ -0,0 +1,126 @@
+"""Claim-vs-evidence judge: verdict classification from claim+evidence text.
+
+Trains the last blocks + a linear head on TinyLiquid's mean-pooled hidden
+state. Class-weighted loss; per-epoch checkpoints; resumable.
+Usage:
+  .venv/bin/python train/train_judge.py --base ckpt/v8_lora/best.pt \
+      --data data/evidence_judge.jsonl --ckpt ckpt/judge --epochs 20
+"""
+import argparse, json, random, 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
+
+VALS = ["true statement", "false statement", "supports", "refutes", "not_enough_info"]
+
+def encode_batch(model, items, tok, max_len=192, grad=True):
+    xs, ps = [], []
+    for it in items:
+        xs.append(tok.encode(it["user"]).ids[:max_len])
+        ps.append(1 if it.get("persona", "analyst") != "skeptic" else 2)
+    L = max(len(x) for x in xs)
+    buf = torch.zeros(len(xs), L, dtype=torch.long)
+    for j, x in enumerate(xs):
+        buf[j, :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)
+    h = h * mask.unsqueeze(-1)
+    return h.sum(1) / mask.sum(1, keepdim=True)
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--base", default="ckpt/v8_lora/best.pt")
+    ap.add_argument("--data", default="data/evidence_judge.jsonl")
+    ap.add_argument("--tok", default="data/tokenizer.json")
+    ap.add_argument("--ckpt", default="ckpt/judge")
+    ap.add_argument("--epochs", type=int, default=20)
+    ap.add_argument("--batch", type=int, default=32)
+    ap.add_argument("--lr", type=float, default=1e-4)
+    ap.add_argument("--threads", type=int, default=4)
+    ap.add_argument("--seed", type=int, default=17)
+    args = ap.parse_args()
+    torch.set_num_threads(args.threads)
+    rng = random.Random(args.seed)
+    tok = load_tokenizer(args.tok)
+
+    rows = [json.loads(l) for l in open(args.data, encoding="utf-8") if l.strip()]
+    rows = [r for r in rows if r.get("verdict") in VALS]
+    # stratified split
+    by_v = {v: [r for r in rows if r["verdict"] == v] for v in VALS}
+    tr, va = [], []
+    for v, lst in by_v.items():
+        rng.shuffle(lst)
+        n_va = max(2, int(len(lst) * 0.15))
+        va += lst[:n_va]; tr += lst[n_va:]
+    rng.shuffle(tr); rng.shuffle(va)
+    counts = {v: sum(1 for r in tr if r["verdict"] == v) for v in VALS}
+    w = torch.tensor([sum(counts.values()) / max(1, counts[v]) for v in VALS], dtype=torch.float32)
+    print(f"train {len(tr)} val {len(va)} weights {w.tolist()}", flush=True)
+
+    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():
+        p.requires_grad = True
+    for p in model.norm_out.parameters():
+        p.requires_grad = True
+    head = nn.Linear(cfg.d_model, len(VALS))
+    model.train()
+
+    out = Path(args.ckpt); out.mkdir(parents=True, exist_ok=True)
+    start_ep, best_acc = 0, 0.0
+    ck = out / "judge.pt"
+    if ck.exists() and args.epochs > 0:
+        st = torch.load(ck, map_location="cpu")
+        model.load_state_dict(st["model"])
+        head.load_state_dict(st["head"])
+        start_ep, best_acc = st["epoch"] + 1, st.get("best_acc", 0.0)
+        print(f"resumed at epoch {start_ep} (best_acc {best_acc:.3f})", flush=True)
+
+    params = [p for p in model.parameters() if p.requires_grad] + list(head.parameters())
+    opt = torch.optim.AdamW(params, lr=args.lr, weight_decay=0.01)
+    t0 = time.time()
+    for ep in range(start_ep, args.epochs):
+        rng.shuffle(tr)
+        for i in range(0, len(tr), args.batch):
+            batch = tr[i:i + args.batch]
+            h = encode_batch(model, batch, tok, grad=True)
+            logits = head(h)
+            target = torch.tensor([VALS.index(b["verdict"]) for b in batch])
+            loss = F.cross_entropy(logits, target, weight=w)
+            opt.zero_grad(); loss.backward()
+            torch.nn.utils.clip_grad_norm_(params, 1.0)
+            opt.step()
+        # eval
+        model.eval(); head.eval()
+        with torch.no_grad():
+            h = encode_batch(model, va, tok, grad=False)
+            preds = head(h).argmax(-1).numpy()
+        targets = np.array([VALS.index(b["verdict"]) for b in va])
+        acc = (preds == targets).mean()
+        per = {VALS[c]: round(float((preds == targets)[targets == c].mean()), 3) if (targets == c).any() else 0.0 for c in range(len(VALS))}
+        print(f"epoch {ep+1}/{args.epochs} val_acc {acc:.3f} per-class {per} ({time.time()-t0:.0f}s)", flush=True)
+        t0 = time.time()
+        if acc > best_acc:
+            best_acc = acc
+        torch.save({"model": model.state_dict(), "head": head.state_dict(), "epoch": ep,
+                    "best_acc": best_acc, "config": cfg.__dict__, "vals": VALS, "base": args.base},
+                   out / "judge.pt")
+        model.train(); head.train()
+    print(f"done -> {out}/judge.pt best_acc {best_acc:.3f}", flush=True)
+
+if __name__ == "__main__":
+    main()

train/train_lm.py

@@ -0,0 +1,273 @@
+"""Causal LM pretraining for TinyLiquid on packed uint16 token files.
+
+Usage:
+  .venv/bin/python train/train_lm.py --data data/train.bin --val data/valid.bin \
+      --config tiny10m --ckpt ckpt/nlp --steps 12000
+Resume:
+  .venv/bin/python train/train_lm.py ... --resume ckpt/nlp
+"""
+
+import argparse
+import json
+import math
+import os
+import random
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+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, PERSONA_TOKENS
+
+SAVE_PREFIX = "model"
+
+
+def parse_args():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data", default="data/train.bin")
+    ap.add_argument("--val", default="data/valid.bin")
+    ap.add_argument("--tok", default="data/tokenizer.json")
+    ap.add_argument("--config", default="tiny10m")
+    ap.add_argument("--ckpt", default="ckpt/nlp")
+    ap.add_argument("--resume", default=None)
+    ap.add_argument("--init-from", default=None,
+                    help="load model weights only (fresh optimizer + LR schedule)")
+    ap.add_argument("--batch", type=int, default=16)
+    ap.add_argument("--seq", type=int, default=256)
+    ap.add_argument("--lr", type=float, default=3e-4)
+    ap.add_argument("--min-lr", type=float, default=1e-5)
+    ap.add_argument("--warmup", type=int, default=200)
+    ap.add_argument("--steps", type=int, default=12000)
+    ap.add_argument("--total-steps", type=int, default=None,
+                    help="absolute final step when resuming; overrides additive --steps")
+    ap.add_argument("--log-every", type=int, default=25)
+    ap.add_argument("--eval-every", type=int, default=500)
+    ap.add_argument("--save-every", type=int, default=1000)
+    ap.add_argument("--val-batches", type=int, default=40)
+    ap.add_argument("--sample-max-new", type=int, default=80)
+    ap.add_argument("--seed", type=int, default=42)
+    ap.add_argument("--threads", type=int, default=8)
+    ap.add_argument("--bf16", action="store_true", help="CPU autocast bf16 (SVE2/BF16 ARMv9)")
+    ap.add_argument("--nan-rollback", type=int, default=50,
+                    help="auto-resume from last checkpoint after N consecutive non-finite steps (0=off)")
+    ap.add_argument("--mtp", type=int, default=0,
+                    help="multi-token prediction aux heads (Meta MTP; 0=off)")
+    return ap.parse_args()
+
+
+def load_bin(path, memmap_threshold=300 * 1024 * 1024):
+    size = Path(path).stat().st_size
+    if size >= memmap_threshold:
+        # memory-map large corpora; sample_batch converts slices on the fly
+        return np.memmap(path, dtype=np.uint16, mode="r")
+    arr = np.fromfile(path, dtype=np.uint16)
+    return torch.from_numpy(arr.astype(np.int64))
+
+
+def make_batches(tokens: torch.Tensor, batch: int, seq: int, rng: random.Random):
+    # non-overlapping windows; cycle through them in order
+    n = (len(tokens) - 1) // seq
+    if n <= 0:
+        raise ValueError("corpus too small for seq")
+    return torch.arange(0, n * seq, seq, dtype=torch.long)
+
+
+def sample_batch(tokens: torch.Tensor, offsets, step: int, batch: int, seq: int):
+    pos = (step * batch) % len(offsets)
+    idx = offsets[pos : pos + batch]
+    if len(idx) < batch:
+        idx = torch.cat([idx, offsets[: batch - len(idx)]])
+    if isinstance(tokens, np.memmap):
+        buf = torch.stack([torch.from_numpy(tokens[int(s): int(s) + seq].astype(np.int64)) for s in idx])
+    else:
+        buf = torch.stack([tokens[int(s): int(s) + seq] for s in idx])
+    return buf[:, :-1], buf[:, 1:]
+
+
+@torch.no_grad()
+def evaluate(model, val_tokens, batch, seq, n_batches, rng, persona=0, bf16=False):
+    model.eval()
+    total, count = 0.0, 0
+    offsets = make_batches(val_tokens, batch, seq, rng)
+    for i in range(n_batches):
+        x, y = sample_batch(val_tokens, offsets, i, batch, seq)
+        p = torch.full((batch,), persona, dtype=torch.long) if persona else None
+        if bf16:
+            with torch.autocast("cpu", dtype=torch.bfloat16):
+                logits = model(x, persona_ids=p)
+                loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1))
+        else:
+            logits = model(x, persona_ids=p)
+            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1))
+        total += loss.item() * y.numel()
+        count += y.numel()
+    model.train()
+    return total / count
+
+
+@torch.no_grad()
+def sample(model, tok, prompt, max_new, persona=0, temp=0.8, top_k=40):
+    ids = tok.encode(prompt).ids
+    out = model.generate(tok, ids, persona_id=persona, max_new=max_new,
+                         temperature=temp, top_k=top_k)
+    return tok.decode(out[len(ids):])
+
+
+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)
+
+    ckpt_dir = Path(args.ckpt)
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
+
+    tok = load_tokenizer(args.tok)
+    cfg = TinyLiquidConfig(vocab_size=tok.get_vocab_size(), mtp_heads=args.mtp,
+                            **CONFIGS[args.config])
+
+    model = TinyLiquid(cfg)
+    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.1)
+    step, best_val = 0, float("inf")
+    grad_accum = 1  # keep simple on CPU
+    resume_from = 0
+
+    if args.init_from:
+        path = latest_ckpt(args.init_from)
+        if path:
+            sd = torch.load(path, map_location="cpu")
+            if args.mtp and set(sd["model"]) != set(model.state_dict()):
+                model.load_state_dict(sd["model"], strict=False)
+                print(f"init-from {path} (weights only, mtp heads fresh)", flush=True)
+            else:
+                model.load_state_dict(sd["model"])
+                print(f"init-from {path} (weights only, fresh schedule)", flush=True)
+
+    if args.resume:
+        path = latest_ckpt(args.resume)
+        if path:
+            sd = torch.load(path, map_location="cpu")
+            model.load_state_dict(sd["model"])
+            opt.load_state_dict(sd["opt"])
+            step, best_val = sd["step"], sd.get("best_val", float("inf"))
+            resume_from = step
+            print(f"resumed {path} at step {step}", flush=True)
+
+    train_tokens = load_bin(args.data)
+    val_tokens = load_bin(args.val)
+    offsets = make_batches(train_tokens, args.batch, args.seq, rng)
+    # shuffle windows (seed-fixed): every batch is a random domain mix, so a
+    # hard slice of the corpus can never dominate a whole batch
+    offsets = offsets[torch.randperm(len(offsets), generator=torch.Generator().manual_seed(args.seed))]
+
+    model.train()
+    total_params = model.num_params()
+    print(f"params: {total_params/1e6:.2f}M  train_tokens: {len(train_tokens):,}  "
+          f"steps_per_epoch: {len(offsets)//args.batch}", flush=True)
+
+    def save(path, tag=""):
+        payload = {
+            "model": model.state_dict(),
+            "opt": opt.state_dict(),
+            "step": step,
+            "best_val": best_val,
+            "config": cfg.__dict__,
+            "args": vars(args),
+            "tag": tag,
+        }
+        # Never expose a partially written checkpoint to resume logic.
+        tmp = path.with_suffix(path.suffix + ".tmp")
+        torch.save(payload, str(tmp))
+        os.replace(tmp, path)
+
+    prompts = ["Once upon a time,", "The little girl wanted to", "In the dark forest,"]
+    target_steps = args.total_steps if args.total_steps is not None else resume_from + args.steps
+    if target_steps < resume_from:
+        raise ValueError(f"total target step {target_steps} is before resume step {resume_from}")
+    t0 = time.time()
+    nan_streak = 0
+    while step < target_steps:
+        step += 1
+        x, y = sample_batch(train_tokens, offsets, step - 1, args.batch, args.seq)
+        opt.zero_grad(set_to_none=True)
+        def _fwd():
+            if args.mtp:
+                logits, aux = model.forward_mtp(x)
+                loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1))
+                for k, a in enumerate(aux):
+                    off = k + 2
+                    if args.seq > off:
+                        loss = loss + 0.1 * F.cross_entropy(
+                            a[:, :-off].reshape(-1, a.size(-1)), x[:, off:].reshape(-1))
+                return loss
+            logits = model(x)
+            return F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1))
+        if args.bf16:
+            with torch.autocast("cpu", dtype=torch.bfloat16):
+                loss = _fwd()
+        else:
+            loss = _fwd()
+        if not torch.isfinite(loss):
+            # never let a single bad batch poison the weights: skip the optimizer
+            # step entirely so neither weights nor LR schedule are mutated
+            nan_streak += 1
+            print(f"  !! step {step} non-finite loss {loss.item():.3e}; skipping step", flush=True)
+            if args.nan_rollback > 0 and nan_streak >= args.nan_rollback:
+                rollback = latest_ckpt(ckpt_dir)
+                if rollback is not None:
+                    sd = torch.load(rollback, map_location="cpu")
+                    model.load_state_dict(sd["model"])
+                    opt.load_state_dict(sd["opt"])
+                    step, best_val = sd["step"], sd.get("best_val", float("inf"))
+                    nan_streak = 0
+                    print(f"  !! {args.nan_rollback} consecutive non-finite steps; "
+                          f"rolled back to {rollback} at step {step}", flush=True)
+                else:
+                    nan_streak = 0
+                    print("  !! no checkpoint to roll back to; continuing", flush=True)
+            continue
+        nan_streak = 0
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        lr = args.lr if step <= args.warmup else args.min_lr + 0.5 * (args.lr - args.min_lr) * (
+            1 + math.cos(math.pi * (step - args.warmup) / max(1, target_steps - args.warmup)))
+        for g in opt.param_groups:
+            g["lr"] = lr
+        opt.step()
+
+        if step % args.log_every == 0 or step == 1:
+            dt = time.time() - t0
+            toks = args.batch * args.seq
+            print(f"step {step}/{target_steps} loss {loss.item():.4f} lr {lr:.2e} "
+                  f"{toks*args.log_every/dt:.0f} tok/s", flush=True)
+            t0 = time.time()
+
+        if step % args.eval_every == 0:
+            vl = evaluate(model, val_tokens, args.batch, args.seq, args.val_batches, rng, bf16=args.bf16)
+            if vl < best_val:
+                best_val = vl
+                save(ckpt_dir / "model_best.pt", tag="best")
+            save(ckpt_dir / f"{SAVE_PREFIX}_{step}.pt", tag=f"step{step}")
+            print(f"  [eval step {step}] val_loss {vl:.4f} (best {best_val:.4f})", flush=True)
+            for pr in prompts:
+                txt = sample(model, tok, pr, args.sample_max_new)
+                print(f"  [gen] {pr} {txt}", flush=True)
+            model.train()
+
+        if step % args.save_every == 0:
+            save(ckpt_dir / f"{SAVE_PREFIX}_{step}.pt", tag=f"step{step}")
+            print(f"  saved {ckpt_dir}/{SAVE_PREFIX}_{step}.pt", flush=True)
+
+    save(ckpt_dir / f"{SAVE_PREFIX}_{step}.pt", tag="final")
+    print("done", flush=True)
+
+
+if __name__ == "__main__":
+    main()

train/train_lora.py

@@ -0,0 +1,306 @@
+"""LoRA SFT for TinyLiquid, per the small-model adaptation recipe (LoRA paper).
+
+Freezes the pretrained base, trains low-rank adapters on the gated-MLP linears
+plus persona embeddings, with a KL anchor to the base and a TinyStories PPL
+guard. Best checkpoint is selected by masked SFT holdout loss while PPL < guard.
+Saved checkpoints are FOLDED back into standard model keys (no lora_* in the
+state dict), so hf/export_hf.py works unchanged.
+
+Usage:
+  .venv/bin/python train/train_lora.py --base ckpt/nlp --data data/sft_mix_v5.jsonl \
+      --ckpt ckpt/v5_lora --epochs 2 --lr 3e-4 --r 16 --kl 0.05
+"""
+import argparse, json, math, random, 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, CONFIGS
+from model.utils import latest_ckpt
+from model.tiny_liquid import TinyLiquid
+from data.tokenizer import load_tokenizer
+
+
+def resolve_ckpt(path):
+    p = Path(path)
+    if p.is_file():
+        return p
+    ck = latest_ckpt(p)
+    assert ck, f"no checkpoints in {path}"
+    return ck
+
+USER_T, ASST_T, EOT_T = "<|user|>", "<|assistant|>", "<|endoftext|>"
+PERSONA_T = {"analyst": "<|analyst|>", "skeptic": "<|skeptic|>", "spock": "<|analyst|>", "none": ""}
+P_IDS = {"analyst": 1, "skeptic": 2, "spock": 1, "none": 0}
+
+
+class LoRALinear(nn.Module):
+    def __init__(self, base: nn.Linear, r: int, alpha: float, dropout: float):
+        super().__init__()
+        self.base = base
+        for p in base.parameters():
+            p.requires_grad = False
+        out_f, in_f = base.weight.shape
+        self.lora_a = nn.Parameter(torch.empty(in_f, r))
+        self.lora_b = nn.Parameter(torch.zeros(r, out_f))
+        nn.init.kaiming_uniform_(self.lora_a, a=math.sqrt(5))
+        self.dropout = nn.Dropout(dropout)
+        self.scale = alpha / max(1, r)
+
+    def forward(self, x):
+        return self.base(x) + (self.dropout(x) @ self.lora_a @ self.lora_b) * self.scale
+
+
+def wrap_lora(model: TinyLiquid, r: int, alpha: float, dropout: float):
+    wrapped = []
+    for name, mod in list(model.named_modules()):
+        if isinstance(mod, nn.Linear) and not name.endswith("lm_head"):
+            lora = LoRALinear(mod, r, alpha, dropout)
+            parts = name.split(".")
+            parent = model
+            for p in parts[:-1]:
+                parent = parent._modules[p] if isinstance(parent, nn.Module) else getattr(parent, p)
+            parent._modules[parts[-1]] = lora
+            wrapped.append((name, lora))
+    return wrapped
+
+
+def fold_state_dict(sd, wrapped):
+    out = {}
+    for k, v in sd.items():
+        if any(k.startswith(n + ".") and not k.startswith(n + ".base.") for n, _ in wrapped):
+            continue  # lora_a / lora_b
+        matched = False
+        for name, _ in wrapped:
+            if k.startswith(name + ".base."):
+                out[name + "." + k.split(".base.", 1)[1]] = v.clone()
+                matched = True
+                break
+        if not matched:
+            out[k] = v.clone()
+    for name, lora in wrapped:
+        delta = (lora.lora_a @ lora.lora_b).t() * lora.scale
+        out[name + ".weight"] = out[name + ".weight"] + delta.detach()
+    return out
+
+
+def parse_args():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--base", default="ckpt/nlp")
+    ap.add_argument("--resume", default="", help="resume from latest ckpt in this dir")
+    ap.add_argument("--data", default="data/sft_mix_v5.jsonl")
+    ap.add_argument("--tok", default="data/tokenizer.json")
+    ap.add_argument("--ckpt", default="ckpt/v5_lora")
+    ap.add_argument("--val-bin", default="data/valid.bin")
+    ap.add_argument("--replay-bin", default="", help="tokenized bin to mix as fluency replay (raw full-loss items)")
+    ap.add_argument("--replay-ratio", type=float, default=0.5, help="fraction of replay items in the train mixture (0..1)")
+    ap.add_argument("--epochs", type=int, default=2)
+    ap.add_argument("--batch", type=int, default=8)
+    ap.add_argument("--seq", type=int, default=256)
+    ap.add_argument("--lr", type=float, default=3e-4)
+    ap.add_argument("--r", type=int, default=16)
+    ap.add_argument("--alpha", type=float, default=32.0)
+    ap.add_argument("--dropout", type=float, default=0.05)
+    ap.add_argument("--kl", type=float, default=0.05)
+    ap.add_argument("--eval-every", type=int, default=25)
+    ap.add_argument("--log-every", type=int, default=25)
+    ap.add_argument("--ppl-guard", type=float, default=60.0)
+    ap.add_argument("--resume-best-sft", type=float, default=None)
+    ap.add_argument("--resume-best-ppl", type=float, default=None)
+    ap.add_argument("--val-batches", type=int, default=2)
+    ap.add_argument("--seed", type=int, default=17)
+    ap.add_argument("--threads", type=int, default=4)
+    return ap.parse_args()
+
+
+def tokenize_example(tok, ex, seq, u_id, a_id, eot_id):
+    if "raw" in ex:
+        ids = tok.encode(ex["raw"]).ids + [eot_id]
+        x = torch.tensor(ids[:-1], dtype=torch.long)
+        y = torch.tensor(ids[1:], dtype=torch.long)
+        return x[:seq], y[:seq], torch.ones_like(y[:seq], dtype=torch.bool), 0
+    persona_name = ex.get("persona", "analyst")
+    persona = PERSONA_T.get(persona_name, PERSONA_T["analyst"])
+    p_id = P_IDS.get(persona_name, 1)
+    p_ids = tok.encode(persona).ids if persona else []
+    ids = p_ids + [u_id] + tok.encode(ex["user"]).ids + [a_id] + tok.encode(ex["assistant"]).ids + [eot_id]
+    if len(ids) > seq:
+        return None
+    asst_start = len(p_ids) + 1 + len(tok.encode(ex["user"]).ids) + 1
+    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
+    if int(mask.sum()) < 16:
+        return None
+    return x, y, 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)
+
+
+@torch.no_grad()
+def val_ppl(model, val_bin, batch=4, seq=64, n_batches=2, seed=0):
+    mm = np.memmap(val_bin, dtype=np.uint16, mode="r")
+    total, cnt = 0.0, 0
+    rng = np.random.RandomState(seed)
+    n = (len(mm) - 1) // seq
+    for _ in range(n_batches):
+        s = int(rng.randint(0, n - batch))
+        buf = torch.stack([torch.from_numpy(mm[s * seq + i * seq: s * seq + i * seq + seq].astype(np.int64)) for i in range(batch)])
+        x, y = buf[:, :-1], buf[:, 1:]
+        loss = F.cross_entropy(model(x).reshape(-1, model.cfg.vocab_size), y.reshape(-1))
+        total += loss.item() * y.numel(); cnt += y.numel()
+    return float(np.exp(total / cnt))
+
+
+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, a_id, eot_id = tok.token_to_id(USER_T), tok.token_to_id(ASST_T), tok.token_to_id(EOT_T)
+    assert None not in (u_id, a_id, eot_id)
+
+    raw = [json.loads(l) for l in open(args.data, encoding="utf-8") if l.strip()]
+    teacher_path = resolve_ckpt(args.base)
+    model_path = resolve_ckpt(args.resume) if args.resume else teacher_path
+    resume_ck = torch.load(model_path, map_location="cpu") if args.resume else None
+    base = torch.load(model_path, 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"})
+    cfg.mtp_heads = 0  # MTP is pretrain-only; post-training has no MTP heads
+    model = TinyLiquid(cfg); model.load_state_dict(base["model"], strict=False)
+    teacher_sd = torch.load(teacher_path, map_location="cpu")["model"]
+
+    teacher = None
+    if args.kl > 0:
+        teacher = TinyLiquid(cfg); teacher.load_state_dict(teacher_sd, strict=False); teacher.eval()
+        for p in teacher.parameters(): p.requires_grad = False
+
+    wrapped = wrap_lora(model, args.r, args.alpha, args.dropout)
+    for p in model.parameters():
+        p.requires_grad = False
+    for p in model.persona_emb.parameters():
+        p.requires_grad = True
+    for _, lora in wrapped:
+        lora.base.weight.requires_grad = False
+        lora.lora_a.requires_grad = True
+        lora.lora_b.requires_grad = True
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad],
+                            lr=args.lr, betas=(0.9, 0.95), weight_decay=0.02)
+    print(f"base {model_path.name} | lora adapters {len(wrapped)} | trainable {trainable:,}", flush=True)
+
+    items_all = [tokenize_example(tok, e, args.seq, u_id, a_id, eot_id) for e in raw]
+    items_all = [i for i in items_all if i is not None]
+    rng.shuffle(items_all)
+
+    n_eval = min(128, max(16, len(items_all) // 12))
+    eval_items, train_items = items_all[:n_eval], items_all[n_eval:]
+    if args.replay_bin:
+        mm = np.memmap(args.replay_bin, dtype=np.uint16, mode="r")
+        n = (len(mm) - 1) // args.seq
+        gold_n = max(1, len(train_items))
+        replay_n = int(gold_n * args.replay_ratio / max(1e-9, 1.0 - args.replay_ratio))
+        rr = np.random.RandomState(args.seed + 1)
+        for _ in range(replay_n):
+            s = int(rr.randint(0, n))
+            w = torch.from_numpy(mm[s * args.seq: (s + 1) * args.seq].astype(np.int64))
+            x, y = w[:-1], w[1:]
+            train_items.append((x, y, torch.ones_like(y, dtype=torch.bool), 0))
+        print(f"replay: {replay_n} raw items from {args.replay_bin} (mixture ratio {args.replay_ratio:.2f})", flush=True)
+    print(f"train {len(train_items)} eval {len(eval_items)} filtered {len(raw) - len(items_all)}", flush=True)
+
+    out = Path(args.ckpt); out.mkdir(parents=True, exist_ok=True)
+    best_score = args.resume_best_sft if args.resume_best_sft is not None else float("inf")
+    best_ppl = args.resume_best_ppl if args.resume_best_ppl is not None else float("inf")
+    step = (resume_ck or {}).get("step", 0)
+    start_iter = (resume_ck or {}).get("iter", step)
+    t0 = time.time()
+    total_steps = (len(train_items) // args.batch) * args.epochs
+    if step:
+        print(f"resuming from {model_path} at step {step}/{total_steps} iter {start_iter}", flush=True)
+
+    def run_eval():
+        model.eval()
+        total, n = 0.0, 0
+        for i in range(0, len(eval_items), args.batch):
+            x, y, m, p = collate(eval_items[i:i + args.batch], args.seq)
+            with torch.no_grad():
+                logits = model(x, persona_ids=p).reshape(-1, model.cfg.vocab_size)
+            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()
+        ppl = val_ppl(model, args.val_bin, n_batches=args.val_batches, seed=args.seed + step)
+        model.train()
+        return total / n, ppl
+
+    def save(path, tag=""):
+        sd = fold_state_dict(model.state_dict(), wrapped)
+        torch.save({"model": sd, "step": step, "iter": iter_no, "config": cfg.__dict__, "tag": tag}, str(path))
+
+    model.train()
+    iter_no = 0
+    for ep in range(args.epochs):
+        rng.shuffle(train_items)
+        usable = len(train_items) - len(train_items) % args.batch
+        for i in range(0, usable, args.batch):
+            iter_no += 1
+            if iter_no <= start_iter:
+                continue
+            step += 1
+            x, y, m, p = collate(train_items[i:i + args.batch], args.seq)
+            opt.zero_grad(set_to_none=True)
+            logits = model(x, persona_ids=p)
+            sft_loss = F.cross_entropy(logits.reshape(-1, model.cfg.vocab_size), y.reshape(-1), reduction="none")
+            sft_loss = (sft_loss * m.reshape(-1)).sum() / m.sum()
+            loss = sft_loss
+            if teacher is not None:
+                with torch.no_grad():
+                    t_logits = teacher(x, persona_ids=p)
+                kl = F.kl_div(F.log_softmax(logits.float(), dim=-1), F.softmax(t_logits.float(), dim=-1),
+                              reduction="none").sum(dim=-1)
+                loss = loss + args.kl * (kl * m).sum() / m.sum()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], 0.5)
+            opt.step()
+            if step % args.log_every == 0:
+                print(f"step {step}/{total_steps} loss {loss.item():.4f} sft {sft_loss.item():.4f} "
+                      f"{args.batch * args.seq * args.log_every / max(1e-6, time.time() - t0):.0f} tok/s", flush=True)
+                t0 = time.time()
+            if step % args.eval_every == 0:
+                sft_vl, ppl = run_eval()
+                tag = ""
+                if ppl < best_ppl:
+                    best_ppl = ppl; save(out / "best_ppl.pt", tag="best_ppl"); tag += " [best ppl]"
+                if ppl <= args.ppl_guard and sft_vl < best_score:
+                    best_score = sft_vl; save(out / "best.pt", tag="best"); tag += " [new best]"
+                save(out / f"model_{step}.pt", tag=f"step{step}")
+                try:
+                    sp = tok.encode("<|analyst|><|user|>Find discrepancies between: Account A: The meeting ended at 11am. Account B: The meeting ended at noon.<|assistant|>").ids
+                    with torch.no_grad():
+                        sout = tok.decode(model.generate(tok, sp, persona_id=1, max_new=50, temperature=0.35,
+                                                         top_k=20, repetition_penalty=1.25, no_repeat_ngram_size=4)[len(sp):]).replace("\n", " ").strip()[:180]
+                    print(f"  [eval {step}] sft_val_loss {sft_vl:.4f} val_ppl {ppl:.2f}{tag}", flush=True)
+                    print(f"    sample: {sout}", flush=True)
+                except Exception:
+                    print(f"  [eval {step}] sft_val_loss {sft_vl:.4f} val_ppl {ppl:.2f}{tag}", flush=True)
+    sd = fold_state_dict(model.state_dict(), wrapped)
+    torch.save({"model": sd, "step": step, "config": cfg.__dict__, "tag": "final"}, out / "model_final.pt")
+    print(f"done -> {out} best_sft={best_score:.4f} best_ppl={best_ppl:.2f}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

train/train_sft.py

@@ -0,0 +1,170 @@
+"""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()

train/train_sft2.py

@@ -0,0 +1,191 @@
+"""Guarded SFT for TinyLiquid: chat+SOP+forensic mix with TinyStories retention.
+
+Differs from train_sft.py:
+  * supports raw full-loss retention examples ({"raw": text})
+  * evals BOTH masked SFT holdout loss AND TinyStories val PPL (coherence guard)
+  * keeps best.pt (min sft_val_loss while val_ppl < 90) and best_ppl.pt (min ppl)
+
+Usage:
+  .venv/bin/python train/train_sft2.py --base ckpt/nlp --data data/sft_mix_v2.jsonl \
+      --ckpt ckpt/v2 --epochs 3 --lr 2e-5
+"""
+import argparse, json, math, random, time
+from pathlib import Path
+import numpy as np
+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, ASST_T, EOT_T = "<|user|>", "<|assistant|>", "<|endoftext|>"
+PERSONA_T = {"analyst": "<|analyst|>", "skeptic": "<|skeptic|>", "spock": "<|analyst|>", "none": ""}
+P_IDS = {"analyst": 1, "skeptic": 2, "spock": 1, "none": 0}
+
+def parse_args():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--base", default="ckpt/nlp")
+    ap.add_argument("--resume", default="", help="resume from latest ckpt in this dir")
+    ap.add_argument("--data", default="data/sft_mix_v2.jsonl")
+    ap.add_argument("--tok", default="data/tokenizer.json")
+    ap.add_argument("--ckpt", default="ckpt/v2")
+    ap.add_argument("--val-bin", default="data/valid.bin")
+    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=2e-5)
+    ap.add_argument("--eval-every", type=int, default=25)
+    ap.add_argument("--log-every", type=int, default=25)
+    ap.add_argument("--ppl-guard", type=float, default=90.0)
+    ap.add_argument("--val-batches", type=int, default=2)
+    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):
+    if "raw" in ex:
+        ids = tok.encode(ex["raw"]).ids + [eot_id]
+        x = torch.tensor(ids[:-1], dtype=torch.long)
+        y = torch.tensor(ids[1:], dtype=torch.long)
+        mask = torch.ones_like(y, dtype=torch.bool)
+        return x[:seq], y[:seq], mask[:seq], 0
+    persona = PERSONA_T.get(ex.get("persona", "analyst"), PERSONA_T["analyst"])
+    p_id = P_IDS.get(ex.get("persona"), 1)
+    user_ids = tok.encode(ex["user"]).ids
+    asst_ids = tok.encode(ex["assistant"]).ids
+    if persona:
+        p_ids_ = tok.encode(persona).ids
+        ids = p_ids_ + [u_id] + user_ids + [a_id] + asst_ids + [eot_id]
+        asst_start = len(p_ids_) + 1 + len(user_ids) + 1
+    else:
+        ids = [u_id] + user_ids + [a_id] + asst_ids + [eot_id]
+        asst_start = 1 + len(user_ids) + 1
+    if len(ids) > seq:
+        ids = ids[:seq - 1] + [eot_id]
+    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
+    return x[:seq], y[:seq], mask[:seq], 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)
+
+@torch.no_grad()
+def val_ppl(model, val_bin, batch=4, seq=64, n_batches=2, seed=0):
+    mm = np.memmap(val_bin, dtype=np.uint16, mode="r")
+    total, cnt = 0.0, 0
+    rng = np.random.RandomState(seed)
+    n = (len(mm) - 1) // seq
+    for b in range(n_batches):
+        s = int(rng.randint(0, n - batch))
+        buf = torch.stack([torch.from_numpy(mm[s * seq + i * seq: s * seq + i * seq + seq].astype(np.int64))
+                           for i in range(batch)])
+        x, y = buf[:, :-1], buf[:, 1:]
+        loss = F.cross_entropy(model(x).reshape(-1, 8192), y.reshape(-1))
+        total += loss.item() * y.numel(); cnt += y.numel()
+    return float(np.exp(total / cnt))
+
+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, a_id, eot_id = tok.token_to_id(USER_T), tok.token_to_id(ASST_T), tok.token_to_id(EOT_T)
+    assert None not in (u_id, a_id, eot_id)
+
+    exs = [json.loads(l) for l in open(args.data, encoding="utf-8") if l.strip()]
+    rng.shuffle(exs)
+    n_eval = min(128, max(8, len(exs) // 12))
+    eval_ex, train_ex = exs[:n_eval], exs[n_eval:]
+    print(f"train {len(train_ex)} eval {len(eval_ex)}", flush=True)
+
+    base_path = latest_ckpt(args.resume or args.base)
+    base = torch.load(base_path, map_location="cpu")
+    base_cfg = base.get("config") or CONFIGS["tiny10m"]
+    cfg = TinyLiquidConfig(vocab_size=tok.get_vocab_size(),
+                           **{k: v for k, v in base_cfg.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_path.name}", flush=True)
+
+    out = Path(args.ckpt); out.mkdir(parents=True, exist_ok=True)
+    steps_per_epoch = max(1, len(train_ex) // args.batch)
+    total_steps = steps_per_epoch * args.epochs
+
+    def make_items(exs_):
+        return [tokenize_example(tok, e, args.seq, u_id, a_id, eot_id) for e in exs_]
+
+    eval_items = make_items(eval_ex)
+
+    def run_eval():
+        model.eval()
+        total, n = 0.0, 0
+        for i in range(0, len(eval_items), args.batch):
+            x, y, m, p = collate(eval_items[i:i + args.batch], args.seq)
+            with torch.no_grad():
+                logits = model(x, persona_ids=p).reshape(-1, 8192)
+            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()
+        sft_vl = total / n
+        ppl = val_ppl(model, args.val_bin, n_batches=args.val_batches)
+        model.train()
+        return sft_vl, ppl
+
+    best_guard, best_ppl = float("inf"), float("inf")
+    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).reshape(-1, 8192)
+            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:
+                print(f"step {step}/{total_steps} loss {loss.item():.4f} "
+                      f"{args.batch*args.seq*args.log_every/(time.time()-t0):.0f} tok/s", flush=True)
+                t0 = time.time()
+            if step % args.eval_every == 0:
+                sft_vl, ppl = run_eval()
+                try:
+                    sp = tok.encode("<|analyst|><|user|>Find discrepancies between: Account A: The meeting ended at 11am. Account B: The meeting ended at noon.<|assistant|>").ids
+                    with torch.no_grad():
+                        sout = tok.decode(model.generate(tok, sp, persona_id=1, max_new=50, temperature=0.35,
+                                                         top_k=20, repetition_penalty=1.25,
+                                                         no_repeat_ngram_size=4)[len(sp):]).replace("\n", " ").strip()[:180]
+                    print(f"    sample: {sout}", flush=True)
+                except Exception:
+                    pass
+                tag = ""
+                if ppl < args.ppl_guard and sft_vl < best_guard:
+                    best_guard = sft_vl
+                    torch.save({"model": model.state_dict(), "step": step, "config": cfg.__dict__}, out / "best.pt")
+                    tag += " [new best]"
+                if ppl < best_ppl:
+                    best_ppl = ppl
+                    torch.save({"model": model.state_dict(), "step": step, "config": cfg.__dict__}, out / "best_ppl.pt")
+                    tag += " [best ppl]"
+                torch.save({"model": model.state_dict(), "step": step, "config": cfg.__dict__}, out / f"model_{step}.pt")
+                print(f"  [eval {step}] sft_val_loss {sft_vl:.4f} val_ppl {ppl:.2f}{tag}", flush=True)
+    torch.save({"model": model.state_dict(), "step": step, "config": cfg.__dict__}, out / "model_final.pt")
+    print(f"done -> {out}  best_guard_sft_loss={best_guard:.4f} best_ppl={best_ppl:.2f}", flush=True)
+
+if __name__ == "__main__":
+    main()

train/train_sft_v4.py

@@ -0,0 +1,262 @@
+"""Conservative SFT for TinyLiquid v4.
+
+Designed for the 7.8M TinyLiquid model where full SFT caused catastrophic
+forgetting/model-collapse into broken forensic jargon. Defaults freeze most of
+the model, train only persona embeddings + final liquid block + output norm,
+and add a base-model KL anchor.
+"""
+import argparse
+import json
+import random
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from model.config import CONFIGS, TinyLiquidConfig
+from model.tiny_liquid import TinyLiquid
+from model.utils import latest_ckpt
+from data.tokenizer import load_tokenizer
+
+USER_T, ASST_T, EOT_T = '<|user|>', '<|assistant|>', '<|endoftext|>'
+PERSONA_T = {'analyst': '<|analyst|>', 'skeptic': '<|skeptic|>', 'spock': '<|analyst|>', 'none': ''}
+P_IDS = {'analyst': 1, 'skeptic': 2, 'spock': 1, 'none': 0}
+
+
+def parse_args():
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--base', default='ckpt/nlp')
+    ap.add_argument('--data', default='data/sft_mix_v4.jsonl')
+    ap.add_argument('--tok', default='data/tokenizer.json')
+    ap.add_argument('--ckpt', default='ckpt/v4')
+    ap.add_argument('--val-bin', default='data/valid.bin')
+    ap.add_argument('--epochs', type=int, default=2)
+    ap.add_argument('--batch', type=int, default=8)
+    ap.add_argument('--seq', type=int, default=256)
+    ap.add_argument('--lr', type=float, default=8e-6)
+    ap.add_argument('--kl', type=float, default=0.05)
+    ap.add_argument('--eval-every', type=int, default=25)
+    ap.add_argument('--log-every', type=int, default=25)
+    ap.add_argument('--ppl-guard', type=float, default=45.0)
+    ap.add_argument('--val-batches', type=int, default=2)
+    ap.add_argument('--resume', default=None, help='continue from this checkpoint (model weights + step); teacher/KL anchor stays --base')
+    ap.add_argument('--resume-best-sft', type=float, default=None)
+    ap.add_argument('--resume-best-ppl', type=float, default=None)
+    ap.add_argument('--resume-iter', type=int, default=None, help='iterations already trained (needed when saved step is a shifted label)')
+    ap.add_argument('--train-scope', choices=['head', 'last', 'last2', 'full'], default='last')
+    ap.add_argument('--seed', type=int, default=11)
+    ap.add_argument('--threads', type=int, default=8)
+    return ap.parse_args()
+
+
+def resolve_ckpt(path):
+    p = Path(path)
+    if p.is_file():
+        return p
+    ck = latest_ckpt(p)
+    assert ck, f'no checkpoints in {path}'
+    return ck
+
+
+def tokenize_example(tok, ex, seq, u_id, a_id, eot_id):
+    if 'raw' in ex:
+        ids = tok.encode(ex['raw']).ids + [eot_id]
+        x = torch.tensor(ids[:-1], dtype=torch.long)
+        y = torch.tensor(ids[1:], dtype=torch.long)
+        mask = torch.ones_like(y, dtype=torch.bool)
+        return x[:seq], y[:seq], mask[:seq], 0
+
+    persona_name = ex.get('persona', 'analyst')
+    persona = PERSONA_T.get(persona_name, PERSONA_T['analyst'])
+    p_id = P_IDS.get(persona_name, 1)
+    user_ids = tok.encode(ex['user']).ids
+    asst_ids = tok.encode(ex['assistant']).ids
+    p_ids = tok.encode(persona).ids if persona else []
+    ids = p_ids + [u_id] + user_ids + [a_id] + asst_ids + [eot_id]
+    asst_start = len(p_ids) + 1 + len(user_ids) + 1
+    if len(ids) > seq:
+        return None
+    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
+    if int(mask.sum().item()) < 16:
+        return None
+    return x, y, 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 set_train_scope(model, scope):
+    for p in model.parameters():
+        p.requires_grad = False
+    for p in model.persona_emb.parameters():
+        p.requires_grad = True
+    if scope in {'last', 'last2', 'full'}:
+        for p in model.blocks[-1].parameters():
+            p.requires_grad = True
+        for p in model.norm_out.parameters():
+            p.requires_grad = True
+    if scope in {'last2', 'full'}:
+        for p in model.blocks[-2].parameters():
+            p.requires_grad = True
+    if scope == 'full':
+        for p in model.parameters():
+            p.requires_grad = True
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+@torch.no_grad()
+def val_ppl(model, val_bin, batch=4, seq=64, n_batches=2, seed=0):
+    mm = np.memmap(val_bin, dtype=np.uint16, mode='r')
+    total, cnt = 0.0, 0
+    rng = np.random.RandomState(seed)
+    n = (len(mm) - 1) // seq
+    for _ in range(n_batches):
+        s = int(rng.randint(0, n - batch))
+        buf = torch.stack([torch.from_numpy(mm[s * seq + i * seq: s * seq + i * seq + seq].astype(np.int64)) for i in range(batch)])
+        x, y = buf[:, :-1], buf[:, 1:]
+        loss = F.cross_entropy(model(x).reshape(-1, model.cfg.vocab_size), y.reshape(-1))
+        total += loss.item() * y.numel()
+        cnt += y.numel()
+    return float(np.exp(total / cnt))
+
+
+@torch.no_grad()
+def sample_text(model, tok, prompt, persona_id):
+    ids = tok.encode(prompt).ids
+    out = model.generate(tok, ids, persona_id=persona_id, max_new=60, temperature=0.35, top_k=20, repetition_penalty=1.25, no_repeat_ngram_size=4)
+    return tok.decode(out[len(ids):]).replace('\n', ' ').strip()[:220]
+
+
+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, a_id, eot_id = tok.token_to_id(USER_T), tok.token_to_id(ASST_T), tok.token_to_id(EOT_T)
+    assert None not in (u_id, a_id, eot_id)
+
+    raw_exs = [json.loads(line) for line in open(args.data, encoding='utf-8') if line.strip()]
+    rng.shuffle(raw_exs)
+    base_path = resolve_ckpt(args.base)
+    base = torch.load(base_path, map_location='cpu')
+    resume_path = resolve_ckpt(args.resume) if args.resume else None
+    resume = torch.load(resume_path, map_location='cpu') if resume_path else None
+    config = (resume or 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((resume or base)['model'])
+
+    teacher = None
+    if args.kl > 0:
+        teacher = TinyLiquid(cfg)
+        teacher.load_state_dict(base['model'])
+        teacher.eval()
+        for p in teacher.parameters():
+            p.requires_grad = False
+
+    trainable = set_train_scope(model, args.train_scope)
+    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=args.lr, betas=(0.9, 0.95), weight_decay=0.02)
+    print(f'loaded base {base_path} trainable {trainable}/{sum(p.numel() for p in model.parameters())} scope={args.train_scope}', flush=True)
+
+    items_all = [tokenize_example(tok, e, args.seq, u_id, a_id, eot_id) for e in raw_exs]
+    items_all = [item for item in items_all if item is not None]
+    rng.shuffle(items_all)
+    n_eval = min(128, max(16, len(items_all) // 12))
+    eval_items, train_items = items_all[:n_eval], items_all[n_eval:]
+    print(f'train {len(train_items)} eval {len(eval_items)} filtered {len(raw_exs) - len(items_all)}', flush=True)
+
+    out = Path(args.ckpt)
+    out.mkdir(parents=True, exist_ok=True)
+    best_score = args.resume_best_sft if args.resume_best_sft is not None else float('inf')
+    best_ppl = args.resume_best_ppl if args.resume_best_ppl is not None else float('inf')
+    step = resume.get('step', 0) if resume is not None else 0
+    start_iter = args.resume_iter if args.resume_iter is not None else (resume.get('iter', step) if resume is not None else 0)
+    start_step = step
+    t0 = time.time()
+    total_steps = (len(train_items) // args.batch) * args.epochs
+    if start_step:
+        print(f'resuming from {resume_path} at step {start_step}/{total_steps} iter {start_iter}', flush=True)
+
+    def run_eval():
+        model.eval()
+        total, n = 0.0, 0
+        for i in range(0, len(eval_items), args.batch):
+            x, y, m, p = collate(eval_items[i:i + args.batch], args.seq)
+            logits = model(x, persona_ids=p).reshape(-1, model.cfg.vocab_size)
+            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()
+        ppl = val_ppl(model, args.val_bin, n_batches=args.val_batches, seed=args.seed + step)
+        model.train()
+        return total / n, ppl
+
+    model.train()
+    iter_no = 0
+    for ep in range(args.epochs):
+        rng.shuffle(train_items)
+        usable = len(train_items) - len(train_items) % args.batch
+        for i in range(0, usable, args.batch):
+            iter_no += 1
+            if iter_no <= start_iter:
+                continue
+            step += 1
+            x, y, m, p = collate(train_items[i:i + args.batch], args.seq)
+            opt.zero_grad(set_to_none=True)
+            logits = model(x, persona_ids=p)
+            sft_loss = F.cross_entropy(logits.reshape(-1, model.cfg.vocab_size), y.reshape(-1), reduction='none')
+            sft_loss = (sft_loss * m.reshape(-1)).sum() / m.sum()
+            loss = sft_loss
+            if teacher is not None:
+                with torch.no_grad():
+                    t_logits = teacher(x, persona_ids=p)
+                kl = F.kl_div(
+                    F.log_softmax(logits.float(), dim=-1),
+                    F.softmax(t_logits.float(), dim=-1),
+                    reduction='none',
+                ).sum(dim=-1)
+                kl = (kl * m).sum() / m.sum()
+                loss = loss + args.kl * kl
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], 0.5)
+            opt.step()
+
+            if step % args.log_every == 0:
+                print(f'step {step}/{total_steps} loss {loss.item():.4f} sft {sft_loss.item():.4f} {args.batch * args.seq * args.log_every / max(1e-6, time.time() - t0):.0f} tok/s', flush=True)
+                t0 = time.time()
+            if step % args.eval_every == 0:
+                sft_vl, ppl = run_eval()
+                tag = ''
+                if ppl < best_ppl:
+                    best_ppl = ppl
+                    torch.save({'model': model.state_dict(), 'step': step, 'iter': iter_no, 'config': cfg.__dict__}, out / 'best_ppl.pt')
+                    tag += ' [best ppl]'
+                if ppl <= args.ppl_guard and sft_vl < best_score:
+                    best_score = sft_vl
+                    torch.save({'model': model.state_dict(), 'step': step, 'iter': iter_no, 'config': cfg.__dict__}, out / 'best.pt')
+                    tag += ' [new best]'
+                torch.save({'model': model.state_dict(), 'step': step, 'iter': iter_no, 'config': cfg.__dict__}, out / f'model_{step}.pt')
+                print(f'  [eval {step}] sft_val_loss {sft_vl:.4f} val_ppl {ppl:.2f}{tag}', flush=True)
+                print('  sample:', sample_text(model, tok, '<|analyst|><|user|>Find discrepancies between: Account A: The meeting ended at 11am. Account B: The meeting ended at noon.<|assistant|>', 1), flush=True)
+    torch.save({'model': model.state_dict(), 'step': step, 'iter': iter_no, 'config': cfg.__dict__}, out / 'model_final.pt')
+    print(f'done -> {out} best_sft={best_score:.4f} best_ppl={best_ppl:.2f}', flush=True)
+
+
+if __name__ == '__main__':
+    main()

train/watchdog_50m.sh

@@ -0,0 +1,58 @@
+#!/usr/bin/env bash
+# Self-healing runner for the 50M continue-pretrain.
+# Resume-safe (ckpt every 500) + auto-restart on any death (OOM/kill/reboot-loss).
+# Stops when the latest checkpoint reaches TARGET_STEPS.
+# Launch detached:  setsid nohup ./train/watchdog_50m.sh > /dev/null 2>&1 < /dev/null &
+set -u
+cd "$(dirname "$0")/.."
+export PYTHONPATH=$PWD MALLOC_ARENA_MAX=2
+
+CKPT=ckpt/hybrid50m_pretrain
+LOG=logs/pretrain50m.log
+TARGET_STEPS=5000
+MAX_STEPS_PER_ATTEMPT=4000   # cap per launch; remaining steps are computed from latest ckpt
+
+last_step() {
+  .venv/bin/python - "$CKPT" << 'PYEOF' 2>/dev/null
+import sys, torch
+from pathlib import Path
+from model.utils import latest_ckpt
+p = latest_ckpt(sys.argv[1])
+if p is None:
+    print(0)
+else:
+    sd = torch.load(p, map_location="cpu", weights_only=False)
+    print(sd.get("step", 0))
+PYEOF
+}
+
+for attempt in $(seq 1 999); do
+  step0=$(last_step)
+  if [ "$step0" -ge "$TARGET_STEPS" ]; then
+    echo "target $TARGET_STEPS reached (step $step0); stopping watchdog" >> "$LOG"
+    break
+  fi
+  need=$((TARGET_STEPS - step0))
+  if [ "$need" -gt "$MAX_STEPS_PER_ATTEMPT" ]; then
+    need=$MAX_STEPS_PER_ATTEMPT
+  fi
+  now=$(date "+%Y-%m-%d %H:%M:%S")
+  echo "=== attempt $attempt start $now (step $step0 -> $TARGET_STEPS, $need steps) ===" >> "$LOG"
+  .venv/bin/python -u train/train_lm.py \
+    --data data/train_phase2b.bin --val data/valid.bin \
+    --tok data/tokenizer.json --config hybrid50m \
+    --ckpt "$CKPT" --resume "$CKPT" \
+    --batch 8 --seq 512 --lr 1.5e-4 --warmup 300 --steps "$need" \
+    --eval-every 500 --save-every 500 --val-batches 40 \
+    --bf16 --threads 6 --nan-rollback 50 --log-every 50 \
+    >> "$LOG" 2>&1
+  code=$?
+  echo "=== attempt $attempt exit code $code $(date) ===" >> "$LOG"
+  step=$(last_step)
+  if [ "$step" -ge "$TARGET_STEPS" ]; then
+    echo "target $TARGET_STEPS reached (step $step); stopping watchdog" >> "$LOG"
+    break
+  fi
+  echo "died at step $step; restarting in 30s" >> "$LOG"
+  sleep 30
+done

train/watchdog_dpo_v22.sh

@@ -0,0 +1,72 @@
+#!/usr/bin/env bash
+# Self-healing LoRA-DPO v22 runner with checkpoint resume.
+# Launch detached: setsid nohup ./train/watchdog_dpo_v22.sh >/dev/null 2>&1 </dev/null &
+#   to survive env session kills. Watchdog owns its own session (PPid=1) so it
+# persists across interactive session ends and re-launches killed children.
+set -u
+cd "$(dirname "$0")/.."
+export PYTHONPATH=$PWD MALLOC_ARENA_MAX=2
+CKPT=ckpt/hybrid50m_v22_dpo_full_recovery
+LOG=logs/dpo_v22_full_recovery.log
+RESUME=ckpt/hybrid50m_v22_dpo/model_600_recovery.pt
+RESUME_FLAGS=(--resume-best-ppl 25.76 --reset-optimizer)
+
+# Continue the newest valid recovery checkpoint when the device kills the
+# child. Recovery checkpoints carry best_ppl and RNG state; the folded source
+# archive needs the explicit legacy metadata and fresh optimizer only once.
+latest_step=-1
+for candidate in "$CKPT"/model_*.pt; do
+  [ -f "$candidate" ] || continue
+  step_name=${candidate##*/model_}
+  step_name=${step_name%.pt}
+  case "$step_name" in ''|*[!0-9]*) continue ;; esac
+  if [ "$step_name" -gt "$latest_step" ]; then
+    latest_step=$step_name
+    RESUME="$candidate"
+    RESUME_FLAGS=()
+  fi
+done
+
+if [ -f "$CKPT/dpo_complete.json" ]; then
+  echo "DPO completion manifest already exists; not relaunching" >> "$LOG"
+  exit 0
+fi
+mkdir -p "$CKPT"
+
+for attempt in $(seq 1 999); do
+  echo "=== dpo attempt $attempt $(date '+%Y-%m-%d %H:%M:%S') ===" >> "$LOG"
+  # Resume from model_600.pt if it exists; otherwise start from model_base
+  if [ -f "$RESUME" ]; then
+    echo "resuming from truthful step-600 recovery archive with fresh optimizer" >> "$LOG"
+    ARGS=(--base ckpt/hybrid50m_v22_lora/best.pt --data data/prefs_p_all.jsonl --tok data/tokenizer16k.json
+          --ckpt "$CKPT" --val-bin data/valid16k.bin --epochs 1 --batch 4 --seq 512
+          --lr 1e-7 --beta 0.05
+          --ppl-guard 60.0 --eval-every 100 --save-every 100 --threads 6
+          "${RESUME_FLAGS[@]}" --resume "$RESUME")
+  else
+    ARGS=(--base ckpt/hybrid50m_v22_lora/best.pt --data data/prefs_p_all.jsonl --tok data/tokenizer16k.json
+          --ckpt "$CKPT" --val-bin data/valid16k.bin --epochs 1 --batch 4 --seq 512
+          --lr 4e-6 --beta 0.05 --lora --r 16 --alpha 32.0 --dropout 0.05
+          --ppl-guard 60.0 --eval-every 100 --save-every 100 --threads 6)
+  fi
+  .venv/bin/python -u train/train_dpo.py "${ARGS[@]}" >> "$LOG" 2>&1
+  code=$?
+  echo "=== dpo attempt $attempt exit $code $(date '+%H:%M:%S') ===" >> "$LOG"
+  if [ -f "$CKPT/dpo_complete.json" ]; then
+    echo "DPO COMPLETE: $CKPT/model_final.pt" >> "$LOG"
+    exit 0
+  fi
+  if [ -f "$CKPT/dpo_failed.json" ]; then
+    echo "DPO failure manifest written; preserving checkpoint for triage" >> "$LOG"
+    exit 3
+  fi
+  if [ "$code" -eq 3 ]; then
+    echo "PPL guard breached; preserving last good checkpoint and stopping" >> "$LOG"
+    exit 3
+  fi
+  if [ "$code" -eq 1 ] || [ "$code" -eq 2 ]; then
+    echo "Configuration or runtime error; stopping instead of retrying blindly" >> "$LOG"
+    exit "$code"
+  fi
+  sleep 20
+done

train/watchdog_dpo_v25.sh

@@ -0,0 +1,33 @@
+#!/usr/bin/env bash
+# Self-healing LFM2 DPO v25 runner (beta=5.0, cosine LR, schema-matched prefs).
+# Resume-safe; stops when model_final.pt exists.
+# Launch detached: setsid nohup ./train/watchdog_dpo_v25.sh >/dev/null 2>&1 </dev/null & disown
+set -u
+cd "$(dirname "$0")/.."
+export PYTHONPATH=$PWD MALLOC_ARENA_MAX=2
+
+CKPT=ckpt/hybrid50m_v25_dpo
+LOG=logs/dpo_v25.log
+mkdir -p "$CKPT"
+
+for attempt in $(seq 1 999); do
+  echo "=== dpo attempt $attempt $(date "+%Y-%m-%d %H:%M:%S") ===" >> "$LOG"
+  ARGS=(--base ckpt/hybrid50m_v25_lora/best.pt --data data/prefs_v23.jsonl --tok data/tokenizer16k.json
+        --ckpt "$CKPT" --val-bin data/valid16k.bin --epochs 1 --batch 4 --seq 512
+        --lr 8e-7 --lr-schedule cosine --lr-min 8e-8 --warmup-frac 0.01
+        --beta 5.0 --margin 0.1 --apo-weight 0.2 --threads 6
+        --ppl-guard 60.0 --eval-every 100 --save-every 100
+        --lora --r 16 --alpha 32.0 --dropout 0.05)
+  if [ -d "$CKPT" ] && ls "$CKPT"/*.pt >/dev/null 2>&1; then
+    ARGS+=(--resume "$CKPT")
+  fi
+  .venv/bin/python -u train/train_dpo.py "${ARGS[@]}" >> "$LOG" 2>&1
+  code=$?
+  echo "=== dpo attempt $attempt exit $code $(date "+%H:%M:%S") ===" >> "$LOG"
+  if [ -f "$CKPT/model_final.pt" ]; then
+    echo "DPO COMPLETE: $CKPT/model_final.pt" >> "$LOG"
+    exit 0
+  fi
+  echo "died; retrying in 20s" >> "$LOG"
+  sleep 20
+done

train/watchdog_dpo_v26.sh

@@ -0,0 +1,33 @@
+#!/usr/bin/env bash
+# Self-healing LFM2 DPO v26 runner (beta=5.0, cosine LR, 154 schema-matched prefs).
+# Resume-safe; stops when model_final.pt exists.
+# Launch detached: setsid nohup ./train/watchdog_dpo_v26.sh >/dev/null 2>&1 </dev/null & disown
+set -u
+cd "$(dirname "$0")/.."
+export PYTHONPATH=$PWD MALLOC_ARENA_MAX=2
+
+CKPT=ckpt/hybrid50m_v26_dpo
+LOG=logs/dpo_v26.log
+mkdir -p "$CKPT"
+
+for attempt in $(seq 1 999); do
+  echo "=== dpo attempt $attempt $(date "+%Y-%m-%d %H:%M:%S") ===" >> "$LOG"
+  ARGS=(--base ckpt/hybrid50m_v26_lora/best.pt --data data/prefs_v26.jsonl --tok data/tokenizer16k.json
+        --ckpt "$CKPT" --val-bin data/valid16k.bin --epochs 1 --batch 4 --seq 512
+        --lr 8e-7 --lr-schedule cosine --lr-min 8e-8 --warmup-frac 0.01
+        --beta 5.0 --margin 0.1 --apo-weight 0.2 --threads 6
+        --ppl-guard 60.0 --eval-every 100 --save-every 100
+        --lora --r 16 --alpha 32.0 --dropout 0.05)
+  if [ -d "$CKPT" ] && ls "$CKPT"/*.pt >/dev/null 2>&1; then
+    ARGS+=(--resume "$CKPT")
+  fi
+  .venv/bin/python -u train/train_dpo.py "${ARGS[@]}" >> "$LOG" 2>&1
+  code=$?
+  echo "=== dpo attempt $attempt exit $code $(date "+%H:%M:%S") ===" >> "$LOG"
+  if [ -f "$CKPT/model_final.pt" ]; then
+    echo "DPO COMPLETE: $CKPT/model_final.pt" >> "$LOG"
+    exit 0
+  fi
+  echo "died; retrying in 20s" >> "$LOG"
+  sleep 20
+done

train/watchdog_eval_candidates.sh

@@ -0,0 +1,52 @@
+#!/usr/bin/env bash
+# Candidate battery evals (tiny-model-eval discipline): SFT best, then DPO
+# best_ppl, each on the SAME fixed battery (main + researcher + red-team).
+# Sequential (one heavy torch job at a time), resume-safe per battery:
+# per-battery logs + eval.py --resume-from, so a device kill never re-runs
+# completed probes. Launch detached:
+#   setsid nohup ./train/watchdog_eval_candidates.sh >/dev/null 2>&1 </dev/null & disown
+set -u
+cd "$(dirname "$0")/.."
+export PYTHONPATH=$PWD MALLOC_ARENA_MAX=2
+TOK=data/tokenizer16k.json
+THREADS=6
+mkdir -p logs
+
+CANDIDATES=(
+  "sft_best|ckpt/hybrid50m_v22_lora/best.pt"
+  "dpo_best_ppl|ckpt/hybrid50m_v22_dpo_full_recovery/best_ppl.pt"
+)
+BATTERIES=( "main|data/eval_probes.jsonl" "researcher|data/probes_researcher.jsonl" "redteam|data/eval_redteam.jsonl" )
+
+echo "=== candidate eval watchdog start $(date '+%Y-%m-%d %H:%M:%S') ===" >> logs/eval_v22_candidates.log
+
+for cand in "${CANDIDATES[@]}"; do
+  name=${cand%%|*}
+  ckpt=${cand##*|}
+  for bat in "${BATTERIES[@]}"; do
+    label=${bat%%|*}
+    probes=${bat##*|}
+    marker="logs/eval_cand_${name}_${label}.done"
+    if [ -f "$marker" ]; then
+      echo "=== [$name/$label] already complete, skipping ===" >> logs/eval_v22_candidates.log
+      continue
+    fi
+    runlog="logs/eval_cand_${name}_${label}.log"
+    attempt=0
+    while [ "$attempt" -lt 8 ]; do
+      attempt=$((attempt+1))
+      echo "=== [$name/$label] attempt $attempt $(date '+%Y-%m-%d %H:%M:%S') ===" >> "$runlog"
+      .venv/bin/python -u research/eval.py --ckpt "$ckpt" --tok "$TOK" --threads "$THREADS" \
+          --probes "$probes" --resume-from "$runlog" >> "$runlog" 2>&1
+      code=$?
+      echo "=== [$name/$label] exit $code $(date '+%H:%M:%S') ===" >> "$runlog"
+      if [ "$code" -eq 0 ]; then
+        touch "$marker"
+        break
+      fi
+      sleep 30
+    done
+  done
+done
+
+echo "=== ALL CANDIDATE EVALS COMPLETE $(date '+%Y-%m-%d %H:%M:%S') ===" >> logs/eval_v22_candidates.log

train/watchdog_eval_merges_v25.sh

@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+# Main-battery evals for v25 merge candidates (tiny-model-posttrain: parallel
+# merge -> eval -> select). Sequential, resume-safe per candidate.
+# Launch detached: setsid nohup ./train/watchdog_eval_merges_v25.sh >/dev/null 2>&1 </dev/null & disown
+set -u
+cd "$(dirname "$0")/.."
+export PYTHONPATH=$PWD MALLOC_ARENA_MAX=2
+TOK=data/tokenizer16k.json
+THREADS=6
+mkdir -p logs
+
+CANDIDATES=(
+  "soup|ckpt/hybrid50m_v25_merges/soup.pt"
+  "taskarith|ckpt/hybrid50m_v25_merges/taskarith_l0p5.pt"
+  "ties|ckpt/hybrid50m_v25_merges/ties_k0p2.pt"
+)
+echo "=== merge candidate eval start $(date '+%Y-%m-%d %H:%M:%S') ===" >> logs/eval_v25_merge_candidates.log
+
+for cand in "${CANDIDATES[@]}"; do
+  name=${cand%%|*}
+  ckpt=${cand##*|}
+  marker="logs/eval_cand_merge_${name}_main.done"
+  if [ -f "$marker" ]; then
+    echo "=== [$name/main] already complete, skipping ===" >> logs/eval_v25_merge_candidates.log
+    continue
+  fi
+  runlog="logs/eval_cand_merge_${name}_main.log"
+  attempt=0
+  while [ "$attempt" -lt 8 ]; do
+    attempt=$((attempt+1))
+    echo "=== [$name/main] attempt $attempt $(date '+%Y-%m-%d %H:%M:%S') ===" >> "$runlog"
+    .venv/bin/python -u research/eval.py --ckpt "$ckpt" --tok "$TOK" --threads "$THREADS" \
+        --probes data/eval_probes.jsonl --resume-from "$runlog" >> "$runlog" 2>&1
+    code=$?
+    echo "=== [$name/main] exit $code $(date '+%H:%M:%S') ===" >> "$runlog"
+    if [ "$code" -eq 0 ]; then
+      touch "$marker"
+      break
+    fi
+    sleep 30
+  done
+done
+
+echo "=== ALL MERGE CANDIDATE EVALS COMPLETE $(date '+%Y-%m-%d %H:%M:%S') ===" >> logs/eval_v25_merge_candidates.log

train/watchdog_eval_v22.sh

@@ -0,0 +1,18 @@
+#!/usr/bin/env bash
+set -u
+cd "$(dirname "$0")/.."
+export PYTHONPATH=$PWD MALLOC_ARENA_MAX=2
+CKPT=ckpt/hybrid50m_v22_dpo_full_recovery
+LOG=logs/eval_v22_detached.log
+if [ -f "$CKPT/eval_complete.json" ]; then exit 0; fi
+echo "=== eval start $(date '+%Y-%m-%d %H:%M:%S') ===" >> "$LOG"
+for pair in "main data/eval_probes.jsonl" "researcher data/probes_researcher.jsonl" "redteam data/eval_redteam.jsonl"; do
+  set -- $pair
+  label=$1; probes=$2
+  echo "== $label battery ==" >> "$LOG"
+  .venv/bin/python research/eval.py --ckpt "$CKPT" --tok data/tokenizer16k.json --threads 6 --probes "$probes" >> "$LOG" 2>&1
+  code=$?
+  if [ "$code" -ne 0 ]; then echo "EVAL FAILED label=$label code=$code" >> "$LOG"; exit "$code"; fi
+done
+printf '%s\n' '{"status":"complete","checkpoint":"ckpt/hybrid50m_v22_dpo_full_recovery/model_final.pt","tokenizer":"data/tokenizer16k.json"}' > "$CKPT/eval_complete.json"
+echo "=== eval complete $(date '+%Y-%m-%d %H:%M:%S') ===" >> "$LOG"

train/watchdog_eval_v23.sh

@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+# Candidate battery evals for v23 (tiny-model-eval discipline): SFT best,
+# each on the SAME fixed battery (main + researcher + red-team). Sequential
+# (one heavy torch job at a time), resume-safe per battery. Launch detached:
+#   setsid nohup ./train/watchdog_eval_v23.sh >/dev/null 2>&1 </dev/null & disown
+set -u
+cd "$(dirname "$0")/.."
+export PYTHONPATH=$PWD MALLOC_ARENA_MAX=2
+TOK=data/tokenizer16k.json
+THREADS=6
+mkdir -p logs
+
+CANDIDATES=(
+  "sft_best|ckpt/hybrid50m_v23_lora/best.pt"
+)
+BATTERIES=( "main|data/eval_probes.jsonl" "researcher|data/probes_researcher.jsonl" "redteam|data/eval_redteam.jsonl" )
+
+echo "=== candidate eval watchdog start $(date '+%Y-%m-%d %H:%M:%S') ===" >> logs/eval_v23_candidates.log
+
+for cand in "${CANDIDATES[@]}"; do
+  name=${cand%%|*}
+  ckpt=${cand##*|}
+  for bat in "${BATTERIES[@]}"; do
+    label=${bat%%|*}
+    probes=${bat##*|}
+    marker="logs/eval_cand_${name}_${label}.done"
+    if [ -f "$marker" ]; then
+      echo "=== [$name/$label] already complete, skipping ===" >> logs/eval_v23_candidates.log
+      continue
+    fi
+    runlog="logs/eval_cand_${name}_${label}.log"
+    attempt=0
+    while [ "$attempt" -lt 8 ]; do
+      attempt=$((attempt+1))
+      echo "=== [$name/$label] attempt $attempt $(date '+%Y-%m-%d %H:%M:%S') ===" >> "$runlog"
+      .venv/bin/python -u research/eval.py --ckpt "$ckpt" --tok "$TOK" --threads "$THREADS" \
+          --probes "$probes" --resume-from "$runlog" >> "$runlog" 2>&1
+      code=$?
+      echo "=== [$name/$label] exit $code $(date '+%H:%M:%S') ===" >> "$runlog"
+      if [ "$code" -eq 0 ]; then
+        touch "$marker"
+        break
+      fi
+      sleep 30
+    done
+  done
+done
+
+echo "=== ALL CANDIDATE EVALS COMPLETE $(date '+%Y-%m-%d %H:%M:%S') ===" >> logs/eval_v23_candidates.log

train/watchdog_eval_v25_dpo.sh

@@ -0,0 +1,48 @@
+#!/usr/bin/env bash
+# Battery evals for v25 DPO (tiny-model-eval discipline): DPO model_final on
+# SAME fixed battery (main + researcher + red-team). Sequential, resume-safe.
+# Launch detached: setsid nohup ./train/watchdog_eval_v25_dpo.sh >/dev/null 2>&1 </dev/null & disown
+set -u
+cd "$(dirname "$0")/.."
+export PYTHONPATH=$PWD MALLOC_ARENA_MAX=2
+TOK=data/tokenizer16k.json
+THREADS=6
+mkdir -p logs
+
+CANDIDATES=(
+  "dpo_v25_final|ckpt/hybrid50m_v25_dpo/model_final.pt"
+)
+BATTERIES=( "main|data/eval_probes.jsonl" "researcher|data/probes_researcher.jsonl" "redteam|data/eval_redteam.jsonl" )
+
+echo "=== v25 DPO candidate eval watchdog start $(date '+%Y-%m-%d %H:%M:%S') ===" >> logs/eval_v25_dpo_candidates.log
+
+for cand in "${CANDIDATES[@]}"; do
+  name=${cand%%|*}
+  ckpt=${cand##*|}
+  for bat in "${BATTERIES[@]}"; do
+    label=${bat%%|*}
+    probes=${bat##*|}
+    marker="logs/eval_cand_${name}_${label}.done"
+    if [ -f "$marker" ]; then
+      echo "=== [$name/$label] already complete, skipping ===" >> logs/eval_v25_dpo_candidates.log
+      continue
+    fi
+    runlog="logs/eval_cand_${name}_${label}.log"
+    attempt=0
+    while [ "$attempt" -lt 8 ]; do
+      attempt=$((attempt+1))
+      echo "=== [$name/$label] attempt $attempt $(date '+%Y-%m-%d %H:%M:%S') ===" >> "$runlog"
+      .venv/bin/python -u research/eval.py --ckpt "$ckpt" --tok "$TOK" --threads "$THREADS" \
+          --probes "$probes" --resume-from "$runlog" >> "$runlog" 2>&1
+      code=$?
+      echo "=== [$name/$label] exit $code $(date '+%H:%M:%S') ===" >> "$runlog"
+      if [ "$code" -eq 0 ]; then
+        touch "$marker"
+        break
+      fi
+      sleep 30
+    done
+  done
+done
+
+echo "=== ALL v25 DPO CANDIDATE EVALS COMPLETE $(date '+%Y-%m-%d %H:%M:%S') ===" >> logs/eval_v25_dpo_candidates.log

train/watchdog_hf_backup.sh

@@ -0,0 +1,21 @@
+#!/usr/bin/env bash
+# Self-healing HF backup runner: restarts hf_backup.py until it reports DONE.
+# Launch detached: setsid nohup ./train/watchdog_hf_backup.sh >/dev/null 2>&1 </dev/null & disown
+set -u
+cd "$(dirname "$0")/.."
+export PYTHONPATH=$PWD
+LOG=logs/hf_backup_all.log
+echo "=== backup watchdog start $(date '+%Y-%m-%d %H:%M:%S') ===" >> "$LOG"
+
+for attempt in $(seq 1 999); do
+  TOKEN=$(cat ~/.cache/huggingface/token 2>/dev/null || echo "${HF_TOKEN:-}")
+  env HF_TOKEN="$TOKEN" .venv/bin/python -u hf_backup.py --stage all >> "$LOG" 2>&1
+  code=$?
+  echo "=== backup attempt $attempt exit $code $(date '+%H:%M:%S') ===" >> "$LOG"
+  if [ $code -eq 0 ]; then
+    echo "BACKUP COMPLETE" >> "$LOG"
+    exit 0
+  fi
+  echo "died; retrying in 20s" >> "$LOG"
+  sleep 20
+done

train/watchdog_lora_v22.sh

@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+# Self-healing LoRA SFT v22 runner. Resume-safe; stops when model_final.pt exists.
+# Launch detached:  setsid nohup ./train/watchdog_lora_v22.sh > /dev/null 2>&1 < /dev/null &
+set -u
+cd "$(dirname "$0")/.."
+export PYTHONPATH=$PWD MALLOC_ARENA_MAX=2
+CKPT=ckpt/hybrid50m_v22_lora
+LOG=logs/sft_v22.log
+
+for attempt in $(seq 1 999); do
+  echo "=== lora attempt $attempt $(date "+%Y-%m-%d %H:%M:%S") ===" >> "$LOG"
+  ARGS=(--base ckpt/hybrid50m_v16k_pretrain --data data/sft_v22.jsonl --tok data/tokenizer16k.json
+        --ckpt "$CKPT" --val-bin data/valid16k.bin --replay-bin data/train_phase2b16k.bin
+        --replay-ratio 0.5 --epochs 3 --batch 4 --seq 512 --lr 5e-5 --r 16 --alpha 32.0
+        --dropout 0.05 --kl 0.1 --ppl-guard 60.0 --eval-every 25 --log-every 25
+        --val-batches 4 --threads 6)
+  if [ -d "$CKPT" ] && ls "$CKPT"/*.pt >/dev/null 2>&1; then
+    ARGS+=(--resume "$CKPT")
+  fi
+  .venv/bin/python -u train/train_lora.py "${ARGS[@]}" >> "$LOG" 2>&1
+  code=$?
+  echo "=== lora attempt $attempt exit $code $(date "+%H:%M:%S") ===" >> "$LOG"
+  if [ -f "$CKPT/model_final.pt" ]; then
+    echo "SFT COMPLETE: $CKPT/model_final.pt" >> "$LOG"
+    exit 0
+  fi
+  echo "died; retrying in 20s" >> "$LOG"
+  sleep 20
+done

train/watchdog_lora_v23.sh

@@ -0,0 +1,31 @@
+#!/usr/bin/env bash
+# Self-healing LoRA SFT v23 runner (class-balanced gold, clean 16k base).
+# Resume-safe; stops when model_final.pt exists.
+# Launch detached:  setsid nohup ./train/watchdog_lora_v23.sh >/dev/null 2>&1 </dev/null & disown
+set -u
+cd "$(dirname "$0")/.."
+export PYTHONPATH=$PWD MALLOC_ARENA_MAX=2
+CKPT=ckpt/hybrid50m_v23_lora
+LOG=logs/sft_v23.log
+mkdir -p "$CKPT"
+
+for attempt in $(seq 1 999); do
+  echo "=== lora attempt $attempt $(date "+%Y-%m-%d %H:%M:%S") ===" >> "$LOG"
+  ARGS=(--base ckpt/hybrid50m_v16k_pretrain --data data/sft_v23.jsonl --tok data/tokenizer16k.json
+        --ckpt "$CKPT" --val-bin data/valid16k.bin --replay-bin data/train_phase2b16k.bin
+        --replay-ratio 0.5 --epochs 3 --batch 4 --seq 512 --lr 5e-5 --r 16 --alpha 32.0
+        --dropout 0.05 --kl 0.1 --ppl-guard 60.0 --eval-every 25 --log-every 25
+        --val-batches 4 --threads 6)
+  if [ -d "$CKPT" ] && ls "$CKPT"/*.pt >/dev/null 2>&1; then
+    ARGS+=(--resume "$CKPT")
+  fi
+  .venv/bin/python -u train/train_lora.py "${ARGS[@]}" >> "$LOG" 2>&1
+  code=$?
+  echo "=== lora attempt $attempt exit $code $(date "+%H:%M:%S") ===" >> "$LOG"
+  if [ -f "$CKPT/model_final.pt" ]; then
+    echo "SFT COMPLETE: $CKPT/model_final.pt" >> "$LOG"
+    exit 0
+  fi
+  echo "died; retrying in 20s" >> "$LOG"
+  sleep 20
+done

train/watchdog_lora_v25.sh

@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+# Self-healing LoRA SFT v25 runner (handcrafted gold 3000-row milestone).
+# Resume-safe; stops when model_final.pt exists.
+# Launch detached: setsid nohup ./train/watchdog_lora_v25.sh >/dev/null 2>&1 </dev/null & disown
+set -u
+cd "$(dirname "$0")/.."
+export PYTHONPATH=$PWD MALLOC_ARENA_MAX=2
+
+CKPT=ckpt/hybrid50m_v25_lora
+LOG=logs/sft_v25.log
+mkdir -p "$CKPT"
+
+for attempt in $(seq 1 999); do
+  echo "=== lora attempt $attempt $(date "+%Y-%m-%d %H:%M:%S") ===" >> "$LOG"
+  ARGS=(--base ckpt/hybrid50m_v16k_pretrain --data data/sft_v25.jsonl --tok data/tokenizer16k.json
+        --ckpt "$CKPT" --val-bin data/valid16k.bin --replay-bin data/train_phase2b16k.bin
+        --replay-ratio 0.5 --epochs 3 --batch 4 --seq 512 --lr 5e-5 --r 16 --alpha 32.0
+        --dropout 0.05 --kl 0.1 --ppl-guard 60.0 --eval-every 25 --log-every 25
+        --val-batches 4 --threads 6)
+  if [ -d "$CKPT" ] && ls "$CKPT"/*.pt >/dev/null 2>&1; then
+    ARGS+=(--resume "$CKPT")
+  fi
+  .venv/bin/python -u train/train_lora.py "${ARGS[@]}" >> "$LOG" 2>&1
+  code=$?
+  echo "=== lora attempt $attempt exit $code $(date "+%H:%M:%S") ===" >> "$LOG"
+  if [ -f "$CKPT/model_final.pt" ]; then
+    echo "SFT COMPLETE: $CKPT/model_final.pt" >> "$LOG"
+    exit 0
+  fi
+  echo "died; retrying in 20s" >> "$LOG"
+  sleep 20
+done

train/watchdog_lora_v26.sh

@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+# Self-healing LoRA SFT v26 runner (handcrafted gold 3000-row milestone).
+# Resume-safe; stops when model_final.pt exists.
+# Launch detached: setsid nohup ./train/watchdog_lora_v26.sh >/dev/null 2>&1 </dev/null & disown
+set -u
+cd "$(dirname "$0")/.."
+export PYTHONPATH=$PWD MALLOC_ARENA_MAX=2
+
+CKPT=ckpt/hybrid50m_v26_lora
+LOG=logs/sft_v26.log
+mkdir -p "$CKPT"
+
+for attempt in $(seq 1 999); do
+  echo "=== lora attempt $attempt $(date "+%Y-%m-%d %H:%M:%S") ===" >> "$LOG"
+  ARGS=(--base ckpt/hybrid50m_v16k_pretrain --data data/sft_v26.jsonl --tok data/tokenizer16k.json
+        --ckpt "$CKPT" --val-bin data/valid16k.bin --replay-bin data/train_phase2b16k.bin
+        --replay-ratio 0.5 --epochs 3 --batch 4 --seq 512 --lr 5e-5 --r 16 --alpha 32.0
+        --dropout 0.05 --kl 0.1 --ppl-guard 60.0 --eval-every 25 --log-every 25
+        --val-batches 4 --threads 6)
+  if [ -d "$CKPT" ] && ls "$CKPT"/*.pt >/dev/null 2>&1; then
+    ARGS+=(--resume "$CKPT")
+  fi
+  .venv/bin/python -u train/train_lora.py "${ARGS[@]}" >> "$LOG" 2>&1
+  code=$?
+  echo "=== lora attempt $attempt exit $code $(date "+%H:%M:%S") ===" >> "$LOG"
+  if [ -f "$CKPT/model_final.pt" ]; then
+    echo "SFT COMPLETE: $CKPT/model_final.pt" >> "$LOG"
+    exit 0
+  fi
+  echo "died; retrying in 20s" >> "$LOG"
+  sleep 20
+done

train/watchdog_merge_v25.sh

@@ -0,0 +1,40 @@
+#!/usr/bin/env bash
+# Parallel merge recipes on the 16k base (tiny-model-posttrain): soup,
+# task arithmetic, TIES. Resume-safe with .done markers per recipe.
+# Launch detached: setsid nohup ./train/watchdog_merge_v25.sh >/dev/null 2>&1 </dev/null & disown
+set -u
+cd "$(dirname "$0")/.."
+export PYTHONPATH=$PWD MALLOC_ARENA_MAX=2
+LOG=logs/merge_v25.log
+OUT=ckpt/hybrid50m_v25_merges
+mkdir -p "$OUT" logs
+echo "=== merge watchdog start $(date '+%Y-%m-%d %H:%M:%S') ===" >> "$LOG"
+
+BASE=ckpt/hybrid50m_v16k_pretrain/model_5000.pt
+SFT=ckpt/hybrid50m_v25_lora/best.pt
+DPO=ckpt/hybrid50m_v25_dpo/model_final.pt
+
+run_recipe() {
+  local name=$1; shift
+  if [ -f "$OUT/$name.pt.done" ] || [ -f "$OUT/$name.pt" ]; then
+    echo "[$name] already merged, skipping" >> "$LOG"
+    return
+  fi
+  echo "[$name] starting $(date '+%H:%M:%S')" >> "$LOG"
+  .venv/bin/python -u train/parallel_merges.py --base "$BASE" --tasks "$SFT" "$DPO" \
+      --out-dir "$OUT" --lambda-ta 0.5 --threads 4 >> "$LOG" 2>&1
+  local code=$?
+  echo "[$name] exit $code $(date '+%H:%M:%S')" >> "$LOG"
+  if [ $code -eq 0 ]; then touch "$OUT/$name.pt.done"; fi
+}
+
+# parallel_merges.py writes soup.pt + taskarith_l0p5.pt in one call
+run_recipe "soup"
+# TIES separately (existing script)
+if [ ! -f "$OUT/ties_k0p2.pt" ]; then
+  echo "[ties] starting $(date '+%H:%M:%S')" >> "$LOG"
+  .venv/bin/python -u train/ties_merge.py --base "$BASE" --tasks "$SFT" "$DPO" \
+      --keep 0.2 --out "$OUT/ties_k0p2.pt" --threads 4 >> "$LOG" 2>&1
+  echo "[ties] exit $? $(date '+%H:%M:%S')" >> "$LOG"
+fi
+echo "=== ALL MERGES DONE $(date '+%Y-%m-%d %H:%M:%S') ===" >> "$LOG"

train/watchdog_reencode_full.sh

@@ -0,0 +1,46 @@
+#!/usr/bin/env bash
+# Self-healing runner for the 528M full-corpus 16k re-encode.
+# Loops: rm partial -> re-encode -> verify against the deterministic anchor
+# (2 complete runs produced 11,544,766 non-empty lines / 520,133,183 tokens;
+# empty EOT segments are dropped by reencode.py, so the retrain's 11,545,267
+# line count is NOT the right expectation).
+# Stops when the output file matches BOTH anchors.
+# Launch detached:  setsid nohup ./train/watchdog_reencode_full.sh > /dev/null 2>&1 < /dev/null &
+set -u
+cd "$(dirname "$0")/.."
+export PYTHONPATH=$PWD
+LOG=logs/tokenizer16k.log
+EXPECT=11544766
+TOK_EXPECT=520133183
+
+for attempt in $(seq 1 99); do
+  echo "=== full re-encode attempt $attempt $(date "+%Y-%m-%d %H:%M:%S") ===" >> "$LOG"
+  rm -f data/train_full16k.bin
+  .venv/bin/python data/reencode.py \
+    --in data/train_full.bin --old-tok data/tokenizer.json \
+    --new-tok data/tokenizer16k.json --out data/train_full16k.bin >> "$LOG" 2>&1
+  code=$?
+  echo "=== attempt $attempt exit $code $(date "+%H:%M:%S") ===" >> "$LOG"
+  lines=$(.venv/bin/python -c "
+import numpy as np
+from data.tokenizer import load_tokenizer
+tok = load_tokenizer('data/tokenizer16k.json')
+eot = tok.token_to_id('<|endoftext|>')
+mm = np.memmap('data/train_full16k.bin', dtype='<u2', mode='r')
+n = len(mm)
+cnt = 0
+for start in range(0, n, 4_000_000):
+    cnt += int((mm[start:start + 4_000_000] == eot).sum())
+print(cnt)
+")
+  tokens=$(.venv/bin/python -c "import os; print(os.path.getsize('data/train_full16k.bin') // 2)")
+  echo "attempt $attempt lines=$lines expect=$EXPECT tokens=$tokens expect_tokens=$TOK_EXPECT" >> "$LOG"
+  if [ "$lines" -eq "$EXPECT" ] && [ "$tokens" -eq "$TOK_EXPECT" ]; then
+    echo "full re-encode COMPLETE ($lines lines, $tokens tokens)" >> "$LOG"
+    exit 0
+  fi
+  echo "incomplete (or killed); retrying in 10s" >> "$LOG"
+  sleep 10
+done
+echo "gave up after 99 attempts" >> "$LOG"
+exit 1

tui/analyst.py

@@ -0,0 +1,291 @@
+"""TinyLiquid Analyst - neon terminal UI (stdlib curses, zero deps).
+
+An OpenCode-style agent terminal themed like Parrot OS meets Matrix:
+near-black background, matrix-green neon, pink/blue/purple accents.
+Chat with the analyst/skeptic personas, run SOP-driven cases, search and
+read the library, manage case files.
+
+Usage:
+  .venv/bin/python tui/analyst.py --ckpt ckpt/dpo
+  .venv/bin/python tui/analyst.py --ckpt ckpt/dpo --case mycase.json
+"""
+
+import argparse
+import curses
+import random
+import textwrap
+
+from tui.engine import AnalystEngine
+from tui.cli import handle_cmd
+
+WRAP = 78
+
+# --- neon palette (256-color) ---
+C_BG = 0          # terminal default (black)
+C_GREEN = 82      # matrix green
+C_PINK = 205
+C_BLUE = 39
+C_PURPLE = 141
+C_CYAN = 51
+C_WHITE = 255
+C_DIM = 240
+
+PAIR = {
+    "bg": 1, "green": 2, "pink": 3, "blue": 4, "purple": 5,
+    "cyan": 6, "white": 7, "dim": 8, "user": 9, "model": 10,
+    "skeptic": 11, "system": 12, "error": 13, "header": 14,
+}
+
+
+def setup_colors(stdscr):
+    curses.start_color()
+    curses.use_default_colors()
+    try:
+        if curses.COLORS >= 256:
+            curses.init_pair(PAIR["bg"], C_WHITE, C_BG)
+            curses.init_pair(PAIR["green"], C_GREEN, C_BG)
+            curses.init_pair(PAIR["pink"], C_PINK, C_BG)
+            curses.init_pair(PAIR["blue"], C_BLUE, C_BG)
+            curses.init_pair(PAIR["purple"], C_PURPLE, C_BG)
+            curses.init_pair(PAIR["cyan"], C_CYAN, C_BG)
+            curses.init_pair(PAIR["white"], C_WHITE, C_BG)
+            curses.init_pair(PAIR["dim"], C_DIM, C_BG)
+            curses.init_pair(PAIR["user"], C_CYAN, C_BG)
+            curses.init_pair(PAIR["model"], C_GREEN, C_BG)
+            curses.init_pair(PAIR["skeptic"], C_PINK, C_BG)
+            curses.init_pair(PAIR["system"], C_PURPLE, C_BG)
+            curses.init_pair(PAIR["error"], C_PINK, C_BG)
+            curses.init_pair(PAIR["header"], C_GREEN, C_BG)
+        else:
+            for k in PAIR:
+                curses.init_pair(PAIR[k], curses.COLOR_WHITE, curses.COLOR_BLACK)
+            curses.init_pair(PAIR["green"], curses.COLOR_GREEN, curses.COLOR_BLACK)
+            curses.init_pair(PAIR["pink"], curses.COLOR_RED, curses.COLOR_BLACK)
+            curses.init_pair(PAIR["blue"], curses.COLOR_CYAN, curses.COLOR_BLACK)
+            curses.init_pair(PAIR["purple"], curses.COLOR_MAGENTA, curses.COLOR_BLACK)
+            curses.init_pair(PAIR["cyan"], curses.COLOR_CYAN, curses.COLOR_BLACK)
+            curses.init_pair(PAIR["user"], curses.COLOR_CYAN, curses.COLOR_BLACK)
+            curses.init_pair(PAIR["model"], curses.COLOR_GREEN, curses.COLOR_BLACK)
+            curses.init_pair(PAIR["skeptic"], curses.COLOR_RED, curses.COLOR_BLACK)
+            curses.init_pair(PAIR["system"], curses.COLOR_MAGENTA, curses.COLOR_BLACK)
+            curses.init_pair(PAIR["error"], curses.COLOR_RED, curses.COLOR_BLACK)
+            curses.init_pair(PAIR["header"], curses.COLOR_GREEN, curses.COLOR_BLACK)
+    except curses.error:
+        pass
+
+
+def wrap(text, width):
+    return "\n".join(textwrap.fill(p, width) for p in text.splitlines()) or " "
+
+
+class Rain:
+    """Minimal matrix-rain strip for the header (lightweight, idle-animated)."""
+
+    def __init__(self, width, rows=1):
+        self.width = max(1, width)
+        self.rows = rows
+        self.cols = [random.randint(0, self.width - 1) for _ in range(self.rows)]
+
+    def frame(self):
+        chars = [" "] * self.width
+        for i, c in enumerate(self.cols):
+            if random.random() < 0.75:
+                chars[c] = random.choice("01$#&@!?ABCDEF")
+            self.cols[i] = (c + random.choice([-1, 0, 1])) % self.width
+        return "".join(chars)
+
+
+class TerminalUI:
+    def __init__(self, stdscr, eng):
+        self.stdscr = stdscr
+        self.eng = eng
+        self.lines = []          # list of (kind, text)
+        self.history = []
+        self.hist_idx = -1
+        self.scroll = 0
+        self.input_text = ""
+        self.cursor = 0
+        self.rain = Rain(10)
+        self.status = f"persona: analyst | case: {eng.case.name}"
+
+    def add(self, kind, text):
+        for ln in wrap(text, WRAP).splitlines() or [""]:
+            self.lines.append((kind, ln))
+        self.scroll = 0
+
+    def header(self, wd):
+        title = " TINYLIQUID ANALYST "
+        left = (wd - len(title)) // 2
+        top = "═" * wd
+        line = " " * left + title
+        return top, line[:wd], self.rain.frame()[:wd]
+
+    def draw(self):
+        h, wd = self.stdscr.getmaxyx()
+        top, title, rain = self.header(wd)
+        body_h = max(1, h - 6)
+
+        self.stdscr.erase()
+        self.stdscr.addnstr(0, 0, top[:wd], wd, curses.color_pair(PAIR["purple"]))
+        self.stdscr.addnstr(1, 0, title[:wd], wd, curses.color_pair(PAIR["header"]))
+        self.stdscr.addnstr(2, 0, rain[:wd], wd, curses.color_pair(PAIR["green"]))
+
+        start = len(self.lines) - body_h - self.scroll
+        start = max(0, min(start, max(0, len(self.lines) - body_h)))
+        for i in range(body_h):
+            idx = start + i
+            if idx < len(self.lines):
+                kind, ln = self.lines[idx]
+                pair = {"user": PAIR["user"], "model": PAIR["model"],
+                        "skeptic": PAIR["skeptic"], "system": PAIR["system"],
+                        "error": PAIR["error"]}.get(kind, PAIR["white"])
+                self.stdscr.addnstr(3 + i, 0, ln[:wd], wd, curses.color_pair(pair))
+
+        self.stdscr.addnstr(h - 3, 0, "─" * wd, wd, curses.color_pair(PAIR["purple"]))
+        self.stdscr.addnstr(h - 2, 0, self.status[:wd], wd,
+                            curses.color_pair(PAIR["system"]) | curses.A_REVERSE)
+        prompt = "❯ "
+        self.stdscr.addnstr(h - 1, 0, prompt + self.input_text[:wd - len(prompt) - 1], wd,
+                            curses.color_pair(PAIR["green"]))
+        try:
+            self.stdscr.move(h - 1, min(self.cursor + len(prompt), wd - 1))
+        except curses.error:
+            pass
+        self.stdscr.refresh()
+
+    def run_cmd(self, cmd):
+        if cmd in ("/quit", "/exit", "/q"):
+            return False
+        if cmd == "/clear":
+            self.lines = []
+            self.add("system", "cleared.")
+            return True
+        self.add("user", cmd if cmd.startswith("/") else "you> " + cmd)
+        if not cmd.startswith("/"):
+            # streamed chat: redraw as tokens arrive
+            buf = []
+            self.lines.append(("model", "…"))
+            try:
+                def on_token(tok_id):
+                    buf.append(tok_id)
+                    self.lines[-1] = ("model", self.eng.tok.decode(buf))
+                    self.draw()
+                self.eng.chat(cmd, on_token=on_token)
+            except Exception as e:
+                self.lines[-1] = ("error", f"error: {e}")
+                self.draw()
+                return True
+            self.lines[-1] = ("model", self.eng.tok.decode(buf))
+            self.draw()
+            return True
+        try:
+            out = handle_cmd(self.eng, cmd)
+        except Exception as e:
+            out = f"error: {e}"
+            self.add("error", out)
+            return True
+        kind = "skeptic" if cmd.startswith(("/skeptic",)) else "model"
+        # render JSON reports nicely
+        if out.strip().startswith("{") or out.strip().startswith("["):
+            import json as _json
+            try:
+                data = _json.loads(out)
+                out = _json.dumps(data, indent=2, ensure_ascii=False)
+            except Exception:
+                pass
+        for ln in out.splitlines():
+            self.add(kind, ln)
+        return True
+
+    def loop(self):
+        setup_colors(self.stdscr)
+        curses.curs_set(1)
+        self.stdscr.timeout(80)
+        self.add("system", "TinyLiquid Analyst — on-device forensic research terminal")
+        self.add("dim", "Authorized research/OSINT only. Outputs are decision support, never a verdict.")
+        self.add("dim", "type /help for commands | /quit to exit")
+        while True:
+            self.draw()
+            try:
+                ch = self.stdscr.get_wch()
+            except curses.error:
+                continue
+            if ch in ("\n", "\r", "KEY_ENTER"):
+                cmd = self.input_text.strip()
+                self.input_text, self.cursor = "", 0
+                self.history.append(cmd)
+                self.hist_idx = len(self.history)
+                if cmd and not self.run_cmd(cmd):
+                    break
+                continue
+            if ch in ("\x7f", "KEY_BACKSPACE"):
+                if self.cursor > 0:
+                    self.input_text = self.input_text[:self.cursor - 1] + self.input_text[self.cursor:]
+                    self.cursor -= 1
+                continue
+            if ch == "KEY_DC":
+                self.input_text = self.input_text[:self.cursor] + self.input_text[self.cursor + 1:]
+                continue
+            if ch == "KEY_LEFT":
+                self.cursor = max(0, self.cursor - 1)
+                continue
+            if ch == "KEY_RIGHT":
+                self.cursor = min(len(self.input_text), self.cursor + 1)
+                continue
+            if ch == "KEY_HOME":
+                self.cursor = 0
+                continue
+            if ch == "KEY_END":
+                self.cursor = len(self.input_text)
+                continue
+            if ch in ("KEY_PPAGE",):
+                self.scroll = min(self.scroll + 5, max(0, len(self.lines) - 1))
+                continue
+            if ch in ("KEY_NPAGE",):
+                self.scroll = max(0, self.scroll - 5)
+                continue
+            if ch == "KEY_UP":
+                if self.history and self.hist_idx > 0:
+                    self.hist_idx -= 1
+                    self.input_text = self.history[self.hist_idx]
+                    self.cursor = len(self.input_text)
+                continue
+            if ch == "KEY_DOWN":
+                if self.history and self.hist_idx < len(self.history) - 1:
+                    self.hist_idx += 1
+                    self.input_text = self.history[self.hist_idx]
+                    self.cursor = len(self.input_text)
+                else:
+                    self.input_text, self.hist_idx = "", len(self.history)
+                continue
+            if isinstance(ch, str) and ch.isprintable():
+                self.input_text = self.input_text[:self.cursor] + ch + self.input_text[self.cursor:]
+                self.cursor += 1
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", default="ckpt/v8_lora/best.pt")
+    ap.add_argument("--tok", default="data/tokenizer.json")
+    ap.add_argument("--case", default=None)
+    ap.add_argument("--threads", type=int, default=4)
+    args = ap.parse_args()
+
+    eng = AnalystEngine(ckpt=args.ckpt, tok_path=args.tok, threads=args.threads)
+    if args.case:
+        eng.load_case(args.case)
+
+    def go(stdscr):
+        ui = TerminalUI(stdscr, eng)
+        ui.status = f"persona: {eng.persona} | case: {eng.case.name} | {eng.ckpt}"
+        ui.loop()
+
+    try:
+        curses.wrapper(go)
+    except KeyboardInterrupt:
+        pass
+    print("session ended.")
+
+
+if __name__ == "__main__":
+    main()

tui/cli.py

@@ -0,0 +1,204 @@
+"""TinyLiquid Analyst - headless terminal client (also scripts/SSH-friendly).
+
+Usage:
+  .venv/bin/python tui/cli.py --ckpt ckpt/dpo                     # interactive REPL
+  .venv/bin/python tui/cli.py --ckpt ckpt/dpo --once "Verify: ..."
+  echo "Verify: ..." | .venv/bin/python tui/cli.py --ckpt ckpt/dpo
+"""
+
+import argparse
+import sys
+
+from tui.engine import AnalystEngine, parse_cmd, BANNER
+
+HELP = """commands:
+  <text>                     chat with the current persona
+  /case <sop> <text>         SOP-driven analysis (verdict + confidence + skeptic)
+  /skeptic <text>            skeptic attack on a claim
+  /search <query>            search the library/corpus index
+  /read <key>                read a document from the index
+  /web <topic>               search the live web (news/archive/wiki)
+  /fetch <url>               pull one URL into the library
+  /pull <topic>              search + pull top docs into the library
+  /tor on|off|status         route .onion fetches through local Tor SOCKS
+  /sop [name]                show one or all procedures
+  /persona <analyst|skeptic|none>
+  /rag on|off                retrieval-grounded chat (auto search + context)
+  /memory on|off             multi-turn conversation memory
+  /mem on|off|stats|<claim>  DNA-helix long-term memory (prior-case recall)
+  /agent <task>              run the procedural agent loop (search/read/verdict)
+  /agents <task>             parallel research swarm: 4 agents, merged report
+  /synth [title]             render case ledger into saved research documents
+  /chart <t>|<l>:<v,..>|...  series chart + saved artifact (cross-synthesis)
+  /opinion <claim>           dual-mind fusion: the model's calibrated opinion
+  /research <question>       research-partner loop: pull, verify, opine, reply
+  /journal [name]           full journalism suite -> CaseFile notebook
+  /gaps                      list the case's open questions (gap ledger)
+  /save <name>               save this case
+  /load <name>               load a case
+  /status                    show state
+  /help                      this help
+  /quit                      exit"""
+
+
+def run_once(eng, text):
+    if not text.strip():
+        return
+    if text.startswith("/case"):
+        rest = text[5:].strip()
+        sop = ""
+        if rest and " " in rest[:60]:
+            maybe, body = rest.split(maxsplit=1)
+            if maybe in eng.__class__.__module__:  # noqa - simple guard
+                pass
+            sop = maybe
+            text = body
+        elif rest:
+            text = rest
+        rep = eng.analyze(text, sop or None)
+        print(json_dump(rep))
+    elif text.startswith("/"):
+        print(handle_cmd(eng, text))
+    else:
+        print(eng.chat(text))
+
+
+def json_dump(rep):
+    import json
+    return json.dumps(rep, indent=2, ensure_ascii=False)
+
+
+def handle_cmd(eng, line):
+    cmd, arg = parse_cmd(line) or ("chat", line)
+    if cmd == "chat":
+        return eng.chat(arg)
+    if cmd == "case":
+        sop = ""
+        body = arg
+        if " " in body:
+            maybe, rest = body.split(maxsplit=1)
+            if maybe in {"claim_verification", "cross_source_discrepancy", "pattern_finding",
+                         "timeline_reconstruction", "historical_truth", "politics_analysis",
+                         "dark_web_research", "terminal_control", "source_triage"}:
+                sop, body = maybe, rest
+        return json_dump(eng.analyze(body, sop or None))
+    if cmd == "skeptic":
+        return eng.generate(arg, persona="skeptic")
+    if cmd == "search":
+        return "\n".join(f"{k} (score {s:.2f})" for k, s in eng.search(arg)) or "(no hits)"
+    if cmd == "read":
+        return eng.read(arg)
+    if cmd == "sop":
+        return eng.sop_text(arg) if arg else "\n".join(list_sops_str())
+    if cmd == "persona":
+        if arg in eng.PERSONA_ID if hasattr(eng, "PERSONA_ID") else arg in ("analyst", "skeptic", "none"):
+            eng.persona = arg
+            return f"persona -> {arg}"
+        return "persona: analyst|skeptic|none"
+    if cmd == "save":
+        return "saved " + eng.save_case(arg or "default")
+    if cmd == "load":
+        return eng.load_case(arg or "default")
+    if cmd == "rag":
+        eng.rag = arg not in ("off", "0", "false")
+        return f"retrieval-grounded chat: {'on' if eng.rag else 'off'}"
+    if cmd == "memory":
+        eng.memory = arg not in ("off", "0", "false")
+        return f"multi-turn memory: {'on' if eng.memory else 'off'}"
+    if cmd == "mem":
+        if arg in ("on", "1", "true", "yes"):
+            eng.helix_on = True
+            return "long-term helix memory: on"
+        if arg in ("off", "0", "false", "no"):
+            eng.helix_on = False
+            return "long-term helix memory: off"
+        if arg == "stats":
+            return json_dump(eng.mem_stats())
+        return eng.recall(arg)
+    if cmd == "agent":
+        return json_dump(eng.agent(arg))
+    if cmd == "agents":
+        return json_dump(eng.agents(arg, n=4))
+    if cmd == "synth":
+        return eng.synthesize(arg or None)
+    if cmd == "chart":
+        return eng.chart(arg)
+    if cmd == "opinion":
+        return json_dump(eng.opinion(arg))
+    if cmd == "research":
+        return json_dump(eng.research(arg))
+    if cmd == "journal":
+        return eng.journal(arg or None)
+    if cmd == "gaps":
+        return "\n".join(f"  {g}" for g in eng.gaps())
+    if cmd == "web":
+        hits = eng.web_search(arg, n=8)
+        if not hits:
+            return "(no web hits)"
+        return "\n".join("[{0}] {1}\n   {2}\n   {3}".format(
+            r["source"], r["title"], r["url"], r["snippet"][:100]) for r in hits)
+    if cmd == "fetch":
+        return json_dump(eng.web_fetch(arg))
+    if cmd == "pull":
+        return json_dump(eng.web_pull(arg, n=3))
+    if cmd == "tor":
+        if arg in ("on", "1", "true", "yes"):
+            eng.tor = True
+            return "tor: on"
+        if arg in ("off", "0", "false", "no"):
+            eng.tor = False
+            return "tor: off"
+        return "tor: {0} | {1}".format(eng.tor, eng.tor_status_msg())
+    if cmd == "status":
+        return (f"persona={eng.persona} memory={'on' if eng.memory else 'off'} "
+                f"helix={'on' if getattr(eng, 'helix_on', False) else 'off'} "
+                f"rag={'on' if eng.rag else 'off'} case={eng.case.name} "
+                f"msgs={len(eng.case.chat)} notes={len(eng.case.ledger)} ckpt={eng.ckpt}")
+    if cmd in ("help", "?"):
+        return HELP
+    if cmd in ("quit", "exit"):
+        raise SystemExit(0)
+    return f"unknown command /{cmd}; try /help"
+
+
+def list_sops_str():
+    from research.agent import SOP_ALIASES
+    return [f"  {k:28s} {', '.join(SOP_ALIASES[k][:3])}" for k in SOP_ALIASES]
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", default="ckpt/v8_lora/best.pt")
+    ap.add_argument("--tok", default="data/tokenizer.json")
+    ap.add_argument("--once", default=None)
+    ap.add_argument("--threads", type=int, default=4)
+    args = ap.parse_args()
+
+    eng = AnalystEngine(ckpt=args.ckpt, tok_path=args.tok, threads=args.threads)
+    print(BANNER, flush=True)
+    print(f"model: {eng.ckpt} | persona: analyst | type /help\n", flush=True)
+
+    if args.once:
+        run_once(eng, args.once)
+        return
+    if not sys.stdin.isatty():
+        for line in sys.stdin:
+            line = line.strip()
+            if not line:
+                continue
+            if line in ("/quit", "/exit"):
+                break
+            print(handle_cmd(eng, line), flush=True)
+        return
+    for line in sys.stdin:
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            print(handle_cmd(eng, line), flush=True)
+        except SystemExit:
+            break
+
+
+if __name__ == "__main__":
+    main()

tui/engine.py

@@ -0,0 +1,467 @@
+"""TinyLiquid Analyst terminal engine.
+
+Purpose-built core for the forensic/dark-web research assistant: chat in the
+analyst/skeptic personas with multi-turn memory, optional retrieval-grounded
+chat (RAG over the library), SOP-driven case analysis, an /agent task loop,
+library search, document reading, and persistent case files. Both the TUI
+(tui/analyst.py) and the headless CLI (tui/cli.py) run on this engine.
+
+Guardrails: research/OSINT only. Dark-web actions require an explicit Tor
+proxy and the crawler's stop rules; nothing illegal is ever in scope.
+"""
+
+import json
+import re
+import threading
+from pathlib import Path
+
+import torch
+
+from model.config import TinyLiquidConfig
+from model.tiny_liquid import TinyLiquid
+from model.utils import latest_ckpt
+from data.tokenizer import load_tokenizer
+from research.index import TinyIndex
+from research.structured import analyst_report
+from research.agent import load_sop, run_case, SOP_ALIASES
+from research import orchestrator as orch
+from research import websearch as ws
+from research import workspace as wkspc
+from research.helix import HelixMemory
+from research.user_journal import UserJournal
+import urllib.parse
+
+ROOT = Path(__file__).resolve().parents[1]
+CASE_DIR = ROOT / "cases"
+
+PERSONA_TOK = {"analyst": "<|analyst|>", "skeptic": "<|skeptic|>", "none": ""}
+PERSONA_ID = {"analyst": 1, "skeptic": 2, "none": 0}
+MEMORY_TURNS = 6          # last N turns injected as conversation history
+MAX_HIST_TOK = 768        # cap history tokens inside the context window
+
+BANNER = (
+    "TinyLiquid Analyst - on-device forensic research terminal\n"
+    "Authorized research/OSINT only. Outputs are decision support, never a verdict.\n"
+    "Dark-web actions: only via an explicit Tor proxy, with rate limits and stop rules.\n"
+)
+
+
+class Case:
+    def __init__(self, name="default"):
+        self.name = name
+        self.ledger = []
+        self.chat = []
+
+    def add(self, role, text):
+        self.chat.append({"role": role, "text": text})
+
+    def note(self, text):
+        self.ledger.append(text)
+
+    def save(self, path: Path):
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(json.dumps({"name": self.name, "ledger": self.ledger,
+                                    "chat": self.chat}, indent=2), encoding="utf-8")
+
+    def load(self, path: Path):
+        d = json.loads(path.read_text(encoding="utf-8"))
+        self.name = d.get("name", self.name)
+        self.ledger = d.get("ledger", [])
+        self.chat = d.get("chat", [])
+
+
+class AnalystEngine:
+    def __init__(self, ckpt="ckpt/v8_lora/best.pt", tok_path="data/tokenizer.json",
+                 library="data/library", threads=8):
+        torch.set_num_threads(threads)
+        self.tok = load_tokenizer(tok_path)
+        self.ckpt = latest_ckpt(ckpt)
+        assert self.ckpt, f"no checkpoints in {ckpt}"
+        sd = torch.load(self.ckpt, map_location="cpu")
+        cfg = TinyLiquidConfig(vocab_size=self.tok.get_vocab_size(),
+                               **{k: v for k, v in sd["config"].items() if k != "vocab_size"})
+        self.model = TinyLiquid(cfg)
+        self.model.load_state_dict(sd["model"])
+        self.model.eval()
+        self.persona = "analyst"
+        self.memory = True
+        self.rag = False
+        self.case = Case()
+        self.library = str(Path(library))
+        self.tor = False
+        self.helix_on = True
+        self.memdb = HelixMemory(str(ROOT / "data" / "helix_memory.jsonl"))
+        self.user_journal = UserJournal()
+        self.index = self._build_index(library)
+
+    def _build_index(self, library):
+        idx = TinyIndex()
+        for f in sorted(Path(library).glob("*.txt")):
+            idx.add(f.stem, str(f), f.read_text(encoding="utf-8", errors="ignore"))
+        for d in (ROOT / "corpus" / "raw",):
+            if d.exists():
+                for f in sorted(d.glob("*.txt")):
+                    idx.add(f"{d.name}/{f.stem}", str(f),
+                            f.read_text(encoding="utf-8", errors="ignore"))
+        idx.build()
+        return idx
+
+    # ---- prompt building ----
+    def _rag_context(self, text):
+        if not self.rag:
+            return ""
+        hits = self.index.query(text, k=2)
+        if not hits:
+            return ""
+        parts = ["Library context:"]
+        for key, score in hits:
+            for k, path, doc in self.index.docs:
+                if k == key:
+                    parts.append(f"[{key}] " + doc[:400].replace("\n", " "))
+                    break
+        return "\n".join(parts) + "\n"
+
+    def _history_prompt(self, text):
+        """Multi-turn memory: last N case turns -> a chat prefix."""
+        if not self.memory:
+            return ""
+        turns = self.case.chat[-(MEMORY_TURNS * 2):] if self.case.chat else []
+        parts = []
+        used = 0
+        for t in reversed(turns):
+            seg = t["text"]
+            if used + len(seg) > MAX_HIST_TOK:
+                break
+            parts.append((t["role"], seg))
+            used += len(seg)
+        parts.reverse()
+        out = []
+        for role, seg in parts:
+            if role == "user":
+                out.append("<|user|>")
+            else:
+                out.append("<|assistant|>")
+            out.append(seg)
+        return "".join(out)
+
+    def _build_prompt(self, text):
+        text2 = self.user_journal.context() + text
+        ctx = self._rag_context(text2)
+        hist = self._history_prompt(text2)
+        p = "<|user|>" + (ctx + text2 if ctx else text2) + "<|assistant|>"
+        return hist + p
+
+    # ---- chat / analysis ----
+    def generate(self, text, persona=None, max_new=180, temp=0.6, on_token=None):
+        persona = persona or self.persona
+        p_token = PERSONA_TOK[persona]
+        prompt = p_token + "<|user|>" + text + "<|assistant|>"
+        ids = self.tok.encode(prompt).ids
+        out = self.model.generate(self.tok, ids, persona_id=PERSONA_ID[persona],
+                                  max_new=max_new, temperature=temp, top_k=40,
+                                  repetition_penalty=1.4, no_repeat_ngram_size=4,
+                                  on_token=on_token)
+        return self.tok.decode(out[len(ids):]).strip()
+
+    def chat(self, text, on_token=None):
+        self.case.add("user", text)
+        self.user_journal.note_thread(text)
+        prompt = self._build_prompt(text)
+        ids = self.tok.encode(prompt).ids
+        out = self.model.generate(self.tok, ids, persona_id=PERSONA_ID[self.persona],
+                                  max_new=180, temperature=0.6, top_k=40,
+                                  repetition_penalty=1.4, no_repeat_ngram_size=4,
+                                  on_token=on_token)
+        reply = self.tok.decode(out[len(ids):]).strip()
+        self.case.add("assistant", reply)
+        return reply
+
+    def analyze(self, text, sop=None):
+        sop_text = load_sop(sop, text) if sop else load_sop(None, text)
+        prior = ""
+        if self.helix_on:
+            rec = self.memdb.recall(text, sop_text)
+            if rec:
+                prior = ("\nPRIOR ANALYSIS OF THIS CASE (recalled):\n" + rec["reasoning"]
+                         + "\n-- use as a prior; do not assume it is still correct.\n\n")
+        user = self.user_journal.context() + prior + sop_text + "\n\nTASK: " + text
+        self.case.add("user", "ANALYSIS: " + text)
+        report = analyst_report(self.model, self.tok, user, persona_id=1,
+                                max_scratch=120, max_reason=60)
+        skeptic = self.generate(
+            f"Act as the skeptic. Attack this conclusion:\nClaim: {text}\n"
+            f"Conclusion: {report.get('verdict', '')} {report.get('reasoning', '')}",
+            persona="skeptic", max_new=100)
+        report["skeptic"] = skeptic
+        report["sop"] = sop or "auto"
+        for q in self._open_questions(report.get("reasoning", "")):
+            note = "OPEN: " + q
+            if note not in self.case.ledger:
+                self.case.ledger.append(note)
+        if self.helix_on:
+            self.memdb.write(
+                claim=text,
+                evidence=sop_text if isinstance(sop_text, str) else "",
+                verdict=report.get("verdict", ""),
+                confidence=report.get("confidence", ""),
+                reasoning=report.get("reasoning", ""),
+            )
+        self.case.add("assistant", json.dumps(report, ensure_ascii=False))
+        self.user_journal.note_thread(text)
+        return report
+
+    def recall(self, text):
+        if not self.helix_on:
+            return "long-term memory is off (/mem on)"
+        rec = self.memdb.recall(text)
+        if not rec:
+            return "(no prior record for this case in long-term memory)"
+        return ("prior verdict: {0} ({1})\nreasoning: {2}\nsource text: {3}".format(
+            rec["verdict"], rec["confidence"], rec["reasoning"],
+            (rec.get("claim", "") or "")[:200]))
+
+    def mem_stats(self):
+        return {"helix": self.helix_on, **self.memdb.stats()}
+
+    # ---- gap ledger: my special touch ----
+    _OPEN_KW = ("missing", "what would settle", "what would change", "what is needed",
+                "what is required", "unsupported detail", "not in the record",
+                "no record of", "no document", "outstanding")
+
+    def _open_questions(self, text):
+        """Pull the open threads out of a verdict's reasoning (pure suit logic)."""
+        if not text:
+            return []
+        out = []
+        sents = re.split(r"(?<=[.!?])\s+", text.replace("\n", " ").strip())
+        for s in sents:
+            low = s.lower()
+            if any(k in low for k in self._OPEN_KW):
+                q = s.strip()
+                if q and q not in out:
+                    out.append(q)
+        return out
+
+    def gaps(self):
+        opens = [l for l in self.case.ledger if l.startswith("OPEN:")]
+        return opens or ["(no open questions on this case yet - run /case analyses)"]
+
+    # ---- dual mind: two cognitive minds fused into one opinion ----
+    def _run_minds(self, doc):
+        """Two minds, one model: analyst + skeptic passes with own memory pools."""
+        from research.fusion import run_two_pass
+        return run_two_pass(self.model, self.tok, doc,
+                            memory=self.memdb if self.helix_on else None)
+
+    def _voice(self, instruction, max_new=150):
+        return self.generate(instruction, persona="analyst", max_new=max_new, temp=0.6)
+
+    def opinion(self, text, evidence=None):
+        from research.verify import deterministic_verdict
+        from research.fusion import fuse, opinion_text
+        doc = text if not evidence else "{0}\nEvidence: {1}".format(text, evidence)
+        rv = deterministic_verdict(doc)
+        rule = rv if rv["verdict"] in ("supports", "refutes", "not enough information") else None
+        a, s = self._run_minds(doc)
+        op = fuse(a, s, rule=rule, sources=getattr(self, "_last_sources", ()))
+        op["spoken"] = opinion_text(op)
+        self.case.add("user", "OPINION: " + text)
+        for q in op["open_questions"]:
+            note = "OPEN: " + q
+            if note not in self.case.ledger:
+                self.case.ledger.append(note)
+        self.case.add("assistant", op["spoken"])
+        return op
+
+    def research(self, question, tor=None, pulls=2):
+        """Research-partner loop: search+pull docs -> verify -> two minds -> reply."""
+        from research.verify import deterministic_verdict
+        from research.fusion import fuse, opinion_text
+        tor = self.tor if tor is None else bool(tor)
+
+        if self.helix_on:
+            rec = self.memdb.recall(question, "")
+            if rec:
+                reply = self._voice(
+                    "A prior record exists for this exact case. Restate your "
+                    "assessment of {0} using: verdict {1} ({2}); reasoning {3}".format(
+                        question, rec["verdict"], rec.get("confidence", ""),
+                        (rec.get("reasoning") or "")[:220]))
+                return {"source": "memory", "question": question, "reply": reply,
+                        "opinion": (rec.get("reasoning") or "")[:300],
+                        "gaps": [], "sources": [], "artifact": None}
+
+        pulled = []
+        try:
+            pulled = ws.pull(question, n=pulls, tor=tor, library_dir=self.library)
+            self._reindex()
+        except Exception:
+            pulled = []
+        self._last_sources = [p.get("file", "") for p in pulled if p.get("file")]
+
+        evidence = ""
+        for key, score in self.index.query(question, k=2):
+            for k, path, doc in self.index.docs:
+                if k == key:
+                    evidence += doc[:600].replace("\n", " ") + " "
+                    break
+        doc = question if not evidence.strip() else question + "\nEvidence: " + evidence[:1200]
+
+        rv = deterministic_verdict(doc)
+        rule = rv if rv["verdict"] in ("supports", "refutes", "not enough information") else None
+        a, s = self._run_minds(doc)
+        op = fuse(a, s, rule=rule, sources=self._last_sources)
+        op["spoken"] = opinion_text(op)
+
+        self.case.add("user", "RESEARCH: " + question)
+        for q in op["open_questions"]:
+            note = "OPEN: " + q
+            if note not in self.case.ledger:
+                self.case.ledger.append(note)
+        art = None
+        try:
+            res = wkspc.synthesize(self.case.ledger, title="research_" + question[:40],
+                                   lines=self.case.ledger)
+            art = res[0] if res else None
+        except Exception:
+            art = None
+
+        srcs = ", ".join(self._last_sources[:4]) or "library index"
+        reply = self._voice(
+            "Research outcome for: {0}\n{1}\nSources: {2}\n"
+            "Reply to the user in your own voice: what the record shows, the "
+            "discrepancy you found, your assessment, and what would change it.".format(
+                question, op["spoken"], srcs))
+        self.case.add("assistant", reply)
+        return {"question": question, "reply": reply, "opinion": op["spoken"],
+                "verdict": op["verdict"], "confidence": op["confidence"],
+                "sources": self._last_sources[:6], "gaps": op["open_questions"],
+                "artifact": art}
+
+    def agent(self, task, max_steps=5):
+        """Procedural agent loop: model issues SEARCH/READ/NOTE/VERDICT actions."""
+        sop_text = load_sop(None, task)
+        plan, ledger = run_case(self.model, self.tok, task, self.index,
+                                sop_text, max_steps=max_steps)
+        report = analyst_report(self.model, self.tok, task, persona_id=1,
+                                max_scratch=90, max_reason=50)
+        report["steps"] = plan
+        self.case.add("user", "AGENT TASK: " + task)
+        self.case.add("assistant", json.dumps(report, ensure_ascii=False))
+        self.user_journal.note_thread(text)
+        return report
+
+    def agents(self, task, n=4, max_steps=5):
+        """Parallel research swarm: n agents under distinct angles, merged."""
+        sop_text = load_sop(None, task)
+        lock = threading.Lock()
+        results = orch.run_parallel(self.model, self.tok, task, self.index,
+                                    sop_text, n=n, lock=lock,
+                                    max_steps=max_steps, library=self.library)
+        merged = orch.synthesize(task, results, self._build_index(self.library))
+        merged["analyst"] = analyst_report(self.model, self.tok, task, persona_id=1,
+                                           max_scratch=90, max_reason=50)
+        self.case.add("user", "AGENTS TASK (" + str(n) + " parallel): " + task)
+        self.case.add("assistant", json.dumps(merged, ensure_ascii=False))
+        return merged
+
+    # ---- research tools ----
+    def search(self, query, k=5):
+        return self.index.query(query, k)
+
+    def read(self, key):
+        for k, path, text in self.index.docs:
+            if k == key:
+                return f"<{k} ({path})>\n" + text[:2500]
+        return "(document not found)"
+
+    def sop_text(self, sop):
+        if sop:
+            return load_sop(sop, "")
+        files = sorted((ROOT / "research" / "sop_library").glob("*.md"))
+        return "\n\n".join(f.read_text(encoding="utf-8") for f in files)
+
+    # ---- case files ----
+    def save_case(self, name):
+        self.case.name = name
+        self.case.save(CASE_DIR / f"{name}.json")
+        return str(CASE_DIR / f"{name}.json")
+
+    def load_case(self, name):
+        path = CASE_DIR / f"{name}.json"
+        if not path.exists():
+            return f"no case file: {path}"
+        self.case = Case(name)
+        self.case.load(path)
+        return f"loaded {path} ({len(self.case.chat)} msgs, {len(self.case.ledger)} notes)"
+
+
+    # ---- live web / dark-web retrieval (client hands) ----
+    def _reindex(self):
+        self.index = self._build_index(self.library)
+
+    def tor_status_msg(self):
+        return ws.tor_status()[1]
+
+    def web_search(self, query, n=8):
+        return ws.search_web(query, limit=n)
+
+    def web_fetch(self, url):
+        doc = ws.fetch(url, tor=self.tor)
+        slug = urllib.parse.urlparse(url).path.rstrip("/").rsplit("/", 1)[-1] or "doc"
+        path = ws.save_doc(self.library, slug, doc["title"], doc["content"])
+        self._reindex()
+        doc["file"] = str(path)
+        return doc
+
+    def web_pull(self, query, n=3):
+        res = ws.pull(query, n=n, tor=self.tor, library_dir=self.library)
+        self._reindex()
+        return res
+
+    # ---- research workspace (analytic artifacts: the sandbox) ----
+    def synthesize(self, title=None, theme=None):
+        """Render the case ledger as saved documents (timeline/table/chart/crossref)."""
+        title = (title or self.case.name or "case").strip()
+        series = getattr(self, "_series", None)
+        res = wkspc.synthesize(self.case.ledger, title, series=series,
+                               theme=theme, lines=self.case.ledger)
+        if not res:
+            return ("no dated artifacts yet - run /case analyses, /agent tasks, "
+                    "or /chart data first")
+        path, doc = res
+        return f"saved {path}\n\n" + doc
+
+    def journal(self, arg=None):
+        """Run the full journalism suite over the case + library (CaseFile)."""
+        from research.journalism import (suite_report, docs_from_library,
+                                         claims_from_ledger)
+        name = (arg or self.case.name or "case").strip()
+        docs = docs_from_library(self.library)
+        claims = claims_from_ledger(self.case.ledger)
+        events = []
+        for line in self.case.ledger:
+            m = re.match(r"^(NOTE|EVENT):\s*(\d{4}-\d{1,2}-\d{1,2})\s+(.+)$",
+                         line.strip())
+            if m:
+                events.append({"when": m.group(2), "what": m.group(3),
+                               "source_id": "-"})
+        md = suite_report(name, docs, claims, events=events)
+        return md
+
+    def chart(self, arg):
+        """Parse 'title | label:a,b,c | label:x,y,z' and render/save a series chart."""
+        title, series = wkspc.parse_series_arg(arg)
+        if len(series) < 2:
+            return "need at least two series: / chart <title> | <label>:<v1,v2,..> | <label>:<v2,v3,..>"
+        self._series = series
+        ch = wkspc.render_series(series, title=title)
+        path = wkspc.save_md(title, ch)
+        return f"saved {path}\n\n" + ch
+
+
+def parse_cmd(line: str):
+    if line.startswith("/"):
+        parts = line[1:].split(maxsplit=1)
+        return (parts[0].lower(), parts[1].strip() if len(parts) > 1 else "")
+    return None