"""
LatticeSpike training driver.
Restores the `train_demo.py` entry point referenced by the README (see
PROJECT_IDEA.md, gap #1): it wires a data source into the end-to-end
LatticeSpike model and trains the combined unsupervised objective.
``--data-source`` selects which pipeline feeds it:
synthetic (default) : the synthetic signal generator (synthetic.py).
Calibrates the anomaly threshold and evaluates
AUPRC + lead time each epoch.
llm : real signals extracted from a HuggingFace causal
LM (llm_signals.py) over a labelled JSONL dataset.
Trains over the whole dataset each epoch and dumps
the final RSNNOutput for offline inspection.
Every internal setting of the run is exposed as a command-line flag (model
architecture, data generation/extraction, loss weights, optimization
schedule, calibration, and checkpointing), so a run can be reconfigured
without touching code.
This module is imported as part of the ``Lattice_Spike`` package (its siblings
use ``from .boxes import ...``), so the block below puts the package's parent
directory on ``sys.path`` and then imports the package by name. That makes
python Lattice_Spike/train.py # from the repo's parent dir
python train.py # from inside the package dir
python -m Lattice_Spike.train # as a module
all work the same way.
Examples
--------
# Quick CPU demo with defaults (synthetic data).
python Lattice_Spike/train.py
# Paper-scale boxes, sigmoid surrogate, robust to 15% contamination.
python Lattice_Spike/train.py --d-box 256 --n-neurons 512 \
--surrogate sigmoid --contaminate 0.15 --trim 0.2
# Resume from a checkpoint and train 5 more epochs.
python Lattice_Spike/train.py --resume runs/demo/last.pt --epochs 5
# Train on real LLM signals extracted from a labelled JSONL dataset.
python Lattice_Spike/train.py --data-source llm --limit 100 --epochs 3
"""
from __future__ import annotations
import argparse
import json
import os
import sys
import time
# --- make the package importable however this file is launched --------------
_HERE = os.path.dirname(os.path.abspath(__file__))
_PARENT = os.path.dirname(_HERE)
_PKG = os.path.basename(_HERE) # the package directory name
if _PARENT not in sys.path:
sys.path.insert(0, _PARENT)
import numpy as np
import torch
_pkg = __import__(_PKG) # e.g. import Lattice_Spike
LatticeSpike = _pkg.LatticeSpike
CombinedLoss = _pkg.CombinedLoss
LossWeights = _pkg.LossWeights
auprc = _pkg.auprc
lead_time = _pkg.lead_time
make_batch = __import__(f"{_PKG}.synthetic", fromlist=["make_batch"]).make_batch
structured_token_table = __import__(
f"{_PKG}.synthetic", fromlist=["structured_token_table"]
).structured_token_table
# llm_signals is imported lazily inside train_llm() -- synthetic-only runs
# (the common case) shouldn't pay for loading the LLM-signal-extraction
# pipeline at all.
# Defaults that differ from the synthetic-mode flag defaults below, applied to
# --data-source llm runs unless the user explicitly overrides the flag.
LLM_MODE_DEFAULTS = {
"topk": 64,
"batch_size": 16,
"epochs": 5,
"lr": 1e-3,
"n_neurons": 64,
"box_encoder": "linear",
"out_dir": "output",
}
# Stable loss schema retained so anchor-off runs remain directly comparable
# with historical runs.
LOSS_KEYS = ("loss", "contain", "vol", "anchor", "recon", "mrl", "gamma")
# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def build_parser() -> argparse.ArgumentParser:
p = argparse.ArgumentParser(
description="Train LatticeSpike on synthetic or real LLM-signal data.",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
# -- data source -----------------------------------------------------------
g = p.add_argument_group("data source")
g.add_argument("--data-source", choices=["synthetic", "llm"], default="synthetic",
help="'synthetic' uses the synthetic generator with AUPRC/lead-"
"time eval; 'llm' extracts real signals from a HuggingFace "
"LM over a labelled JSONL dataset (see 'data (LLM signals)' "
"group). Some flag defaults change under --data-source llm "
f"(see LLM_MODE_DEFAULTS: {LLM_MODE_DEFAULTS}).")
# -- model architecture --------------------------------------------------
g = p.add_argument_group("model")
g.add_argument("--d-box", type=int, default=64,
help="box latent dimension d (a box is a 2d vector).")
g.add_argument("--n-neurons", type=int, default=64,
help="number of LIF neurons in the RSNN.")
g.add_argument("--box-encoder", choices=["mlp", "linear"], default="mlp",
help="hidden-state to box encoder architecture.")
g.add_argument("--mrl-k", type=int, default=16,
help="MRL projection width k for the KL trace (channel 3).")
g.add_argument("--window", type=int, default=8,
help="causal window for spike density / anomaly smoothing.")
g.add_argument("--rho-star", type=float, default=0.15,
help="spike-density burst threshold.")
g.add_argument("--surrogate", choices=["atan", "sigmoid"], default="atan",
help="surrogate-gradient spike function.")
g.add_argument("--ablation", choices=["lstm"], default=None,
help="ablation switch: swap the RSNN dynamics layer for the "
"LSTM ablation (see docs/ablation/"
"lstm_ablation_design.md).")
g.add_argument("--lstm-use-volume", action="store_true",
help="with --ablation lstm: append log V_t to the LSTM input "
"(information-matched fair comparison with the RSNN).")
g.add_argument("--lstm-hidden", type=int, default=None,
help="with --ablation lstm: hidden-size override (default: "
"parameter-matched to the RSNN within 5%%).")
# -- synthetic data ------------------------------------------------------
g = p.add_argument_group("data (synthetic generator)")
g.add_argument("--seq-len", type=int, default=24, help="tokens per sequence T.")
g.add_argument("--d-model", type=int, default=64,
help="hidden-state width fed to the box encoder.")
g.add_argument("--num-layers", type=int, default=8,
help="number of LLM layers L in the synthetic KL trace.")
g.add_argument("--vocab", type=int, default=200, help="synthetic vocab size.")
g.add_argument("--n-in", type=int, default=100,
help="number of 'in-box' token ids (< n-in are grounded).")
g.add_argument("--topk", type=int, default=8,
help="top-K next-token candidates per step (channel 6).")
g.add_argument("--batch-size", type=int, default=32, help="sequences per step.")
g.add_argument("--train-hallu-prob", type=float, default=0.0,
help="fraction of TRAIN sequences that are overtly hallucinated "
"(unsupervised training is normally on grounded data).")
g.add_argument("--eval-hallu-prob", type=float, default=0.5,
help="fraction of EVAL sequences that are hallucinated.")
g.add_argument("--contaminate", type=float, default=0.0,
help="fraction of 'grounded' TRAIN sequences that are secretly "
"hallucinated (tests the trimmed reconstruction loss).")
g.add_argument("--onset-lo", type=float, default=0.3,
help="earliest drift onset as a fraction of T.")
g.add_argument("--onset-hi", type=float, default=0.6,
help="latest drift onset as a fraction of T.")
g.add_argument("--lead-gap-lo", type=int, default=3,
help="min tokens from drift onset to irreversibility t*.")
g.add_argument("--lead-gap-hi", type=int, default=7,
help="max tokens from drift onset to irreversibility t*.")
g.add_argument("--structured-tokens", dest="structured_tokens",
action="store_true", default=True,
help="load the structured box-space token table so the "
"concretization channel is meaningful (recommended).")
g.add_argument("--no-structured-tokens", dest="structured_tokens",
action="store_false",
help="keep the model's random frozen token table.")
# -- LLM signals (only used with --data-source llm) -----------------------
g = p.add_argument_group("data (LLM signals; --data-source llm)")
g.add_argument("--llm-data", type=str, default="data/auto_labelled_llama8b.jsonl",
help="path to the labelled JSONL file (relative to the "
"Lattice_Spike package dir, or absolute).")
g.add_argument("--llm-model", type=str, default="Qwen/Qwen2.5-0.5B-Instruct",
help="HuggingFace model ID or local path used for signal "
"extraction (requires the 'transformers' package).")
g.add_argument("--text-field", type=str, default="generated_answer",
help="JSONL field containing the text to extract signals from.")
g.add_argument("--limit", type=int, default=None,
help="max JSONL records to process (default: all).")
# -- loss weights --------------------------------------------------------
g = p.add_argument_group("loss")
g.add_argument("--lam-vol", type=float, default=0.1, help="volume-loss weight.")
g.add_argument("--lam-anchor", type=float, default=0.0,
help="anchored-contraction-loss weight; 0 disables it while "
"retaining diagnostics.")
g.add_argument("--lam-recon", type=float, default=1.0,
help="temporal-reconstruction-loss weight.")
g.add_argument("--lam-mrl", type=float, default=0.1,
help="MRL nesting-loss weight.")
g.add_argument("--lam-join", type=float, default=0.0,
help="volumetric join/intersection (meet) containment-loss "
"weight; pass 0 to disable it.")
g.add_argument("--lam-rank", type=float, default=0.0,
help="supervised hallucinated-vs-grounded margin ranking weight.")
g.add_argument("--rank-margin", type=float, default=0.1,
help="required anomaly-score gap: hallucinated >= grounded + margin.")
g.add_argument("--rank-pair-mode", choices=["sequence", "batch", "both", "none"],
default="both",
help="within-sequence, cross-sequence, or both token pair sets.")
g.add_argument("--rank-max-pairs", type=int, default=4096,
help="max deterministic pairs per sequence pair; 0 uses all pairs.")
g.add_argument("--trim", type=float, default=0.1,
help="trimmed fraction of high-residual tokens in L_recon "
"(robustness to contamination).")
# -- optimization --------------------------------------------------------
g = p.add_argument_group("optimization")
g.add_argument("--optimizer", choices=["adamw", "adam", "sgd"], default="adamw")
g.add_argument("--lr", type=float, default=2e-3, help="learning rate.")
g.add_argument("--weight-decay", type=float, default=0.0)
g.add_argument("--grad-clip", type=float, default=1.0,
help="max global grad norm (0 disables clipping).")
g.add_argument("--epochs", type=int, default=10)
g.add_argument("--steps-per-epoch", type=int, default=20,
help="optimizer steps (fresh synthetic batch) per epoch.")
# -- calibration / eval --------------------------------------------------
g = p.add_argument_group("calibration & eval")
g.add_argument("--calib-z", type=float, default=4.0,
help="anomaly threshold = mean + z*std on grounded data.")
g.add_argument("--eval-batch", type=int, default=256,
help="sequences in the held-out eval batch.")
# -- checkpointing / runtime --------------------------------------------
g = p.add_argument_group("checkpoint & runtime")
g.add_argument("--out-dir", type=str, default="runs/demo",
help="directory for checkpoints and config.json.")
g.add_argument("--save-every", type=int, default=0,
help="also snapshot every N epochs (0 = only last/best).")
g.add_argument("--no-save", action="store_true",
help="run without writing any checkpoints (dry run).")
g.add_argument("--resume", type=str, default="",
help="checkpoint path to resume model/optimizer state from.")
g.add_argument("--seed", type=int, default=0)
g.add_argument("--device", type=str, default="cpu",
help="cpu or cuda (or cuda:N).")
g.add_argument("--log-every", type=int, default=1,
help="print a training-loss line every N steps.")
g.add_argument("--viz-samples", type=int, default=6,
help="eval-batch examples to snapshot as decision-tree records "
"per tracked epoch in history.jsonl (0 disables sample "
"snapshots; scalar metrics are always logged).")
g.add_argument("--viz-every", type=int, default=1,
help="snapshot samples every N epochs.")
return p
# ---------------------------------------------------------------------------
# helpers
# ---------------------------------------------------------------------------
def set_seed(seed: int) -> None:
torch.manual_seed(seed)
np.random.seed(seed)
def cli_overrides(parser: argparse.ArgumentParser, argv) -> set[str]:
"""Return argparse dest names explicitly set on the command line."""
argv = sys.argv[1:] if argv is None else list(argv)
option_to_dest = {
opt: action.dest
for action in parser._actions
for opt in action.option_strings
}
seen: set[str] = set()
for token in argv:
if token == "--":
break
if not token.startswith("-"):
continue
option = token.split("=", 1)[0]
dest = option_to_dest.get(option)
if dest:
seen.add(dest)
return seen
def apply_defaults(args, values: dict, overrides: set[str], skip: set[str] = frozenset()) -> None:
"""Set args.<name> = value for each (name, value) not in `skip`/`overrides`."""
for name, value in values.items():
if name in skip or name in overrides or not hasattr(args, name):
continue
setattr(args, name, value)
def apply_checkpoint_args(args, ckpt: dict, overrides: set[str]) -> None:
"""Restore saved run settings on resume unless the user overrode them.
Checkpoints store the original CLI args. Reapplying them before model
construction lets `--resume path --epochs N` work for non-default model
shapes, while still allowing intentional command-line overrides.
"""
saved = ckpt.get("args") or {}
apply_defaults(args, saved, overrides, skip={"resume", "epochs", "no_save"})
def _validate_synthetic_args(args) -> list[str]:
"""Checks specific to the synthetic generator (data_source == 'synthetic')."""
errs = []
if not (0 < args.n_in < args.vocab):
errs.append(f"--n-in must satisfy 0 < n_in < vocab "
f"(got n_in={args.n_in}, vocab={args.vocab}); the generator "
f"needs both in-box (<n_in) and out-of-box (>=n_in) ids.")
if args.topk > args.vocab:
errs.append(f"--topk ({args.topk}) cannot exceed --vocab ({args.vocab}).")
if not (0.0 <= args.onset_lo < args.onset_hi <= 1.0):
errs.append(f"--onset-lo/--onset-hi must satisfy "
f"0 <= onset_lo < onset_hi <= 1 "
f"(got {args.onset_lo}, {args.onset_hi}).")
elif int(args.onset_lo * args.seq_len) >= int(args.onset_hi * args.seq_len):
errs.append(f"--seq-len ({args.seq_len}) is too short for the onset range "
f"[{args.onset_lo}, {args.onset_hi}); the integer onset "
f"bounds collapse to an empty range.")
if args.lead_gap_lo < 1 or args.lead_gap_lo > args.lead_gap_hi:
errs.append(f"--lead-gap-lo/--lead-gap-hi must satisfy "
f"1 <= lead_gap_lo <= lead_gap_hi "
f"(got {args.lead_gap_lo}, {args.lead_gap_hi}).")
for name in ("train_hallu_prob", "eval_hallu_prob", "contaminate"):
v = getattr(args, name)
if not (0.0 <= v <= 1.0):
errs.append(f"--{name.replace('_', '-')} must be in [0, 1] (got {v}).")
for name in ("seq_len", "d_model", "num_layers", "eval_batch", "steps_per_epoch"):
v = getattr(args, name)
if v < 1:
errs.append(f"--{name.replace('_', '-')} must be >= 1 (got {v}).")
if args.viz_samples < 0:
errs.append(f"--viz-samples must be >= 0 (got {args.viz_samples}).")
if args.viz_every < 1:
errs.append(f"--viz-every must be >= 1 (got {args.viz_every}).")
return errs
def _validate_llm_args(args) -> list[str]:
"""Checks specific to the LLM-signal pipeline (data_source == 'llm')."""
errs = []
if args.resume:
errs.append("--resume is not supported with --data-source llm.")
if args.limit is not None and args.limit < 1:
errs.append(f"--limit must be >= 1 (got {args.limit}).")
return errs
def _validate_shared_args(args) -> list[str]:
"""Checks that apply to both data sources (model/loss/optimizer flags)."""
errs = []
if args.topk < 1:
errs.append(f"--topk must be >= 1 (got {args.topk}).")
if not (0.0 <= args.trim < 1.0):
errs.append(f"--trim must be in [0, 1) (got {args.trim}).")
for name in ("d_box", "n_neurons", "mrl_k", "window", "batch_size", "epochs"):
v = getattr(args, name)
if v < 1:
errs.append(f"--{name.replace('_', '-')} must be >= 1 (got {v}).")
if args.lr <= 0:
errs.append(f"--lr must be > 0 (got {args.lr}).")
if args.weight_decay < 0:
errs.append(f"--weight-decay must be >= 0 (got {args.weight_decay}).")
if args.grad_clip < 0:
errs.append(f"--grad-clip must be >= 0 (got {args.grad_clip}).")
if not (0.0 <= args.rho_star <= 1.0):
errs.append(f"--rho-star must be in [0, 1] (got {args.rho_star}).")
if args.calib_z < 0:
errs.append(f"--calib-z must be >= 0 (got {args.calib_z}).")
for name in ("lam_vol", "lam_anchor", "lam_recon", "lam_mrl", "lam_join", "lam_rank"):
v = getattr(args, name)
if v < 0:
errs.append(f"--{name.replace('_', '-')} must be >= 0 (got {v}).")
if args.rank_margin < 0:
errs.append(f"--rank-margin must be >= 0 (got {args.rank_margin}).")
if args.rank_max_pairs < 0:
errs.append(f"--rank-max-pairs must be >= 0 (got {args.rank_max_pairs}).")
if args.lam_rank > 0 and args.rank_pair_mode == "none":
errs.append("--rank-pair-mode cannot be none when --lam-rank is positive.")
if args.save_every < 0:
errs.append(f"--save-every must be >= 0 (got {args.save_every}).")
if args.log_every < 0:
errs.append(f"--log-every must be >= 0 (got {args.log_every}).")
return errs
def validate_args(args) -> list[str]:
"""Return a list of human-readable problems with the flag combination.
Caught early so the user gets a clear message instead of a deep stack trace
from inside the data generator (e.g. n_in == vocab leaves no out-of-box ids
and crashes torch.randint(n_in, vocab)).
"""
mode_errs = (_validate_synthetic_args(args) if args.data_source == "synthetic"
else _validate_llm_args(args))
return mode_errs + _validate_shared_args(args)
def resolve_device(args) -> str:
"""Apply the CUDA fallback once and return the device string."""
device = args.device
if device.startswith("cuda") and not torch.cuda.is_available():
print(f"[warn] {device} requested but CUDA is unavailable; using cpu.")
device = args.device = "cpu"
return device
def build_model(args) -> LatticeSpike:
"""Build the LatticeSpike model from args.d_model/num_layers/vocab/etc.
The LLM pipeline (train_llm) doesn't know its real d_model/num_layers/
vocab_size until the HuggingFace model is loaded, so it sets those
args fields (and args.structured_tokens = False) before calling this,
the same way resolve_device() sets args.device in place.
"""
model = LatticeSpike(
d_model=args.d_model,
vocab_size=args.vocab,
d_box=args.d_box,
num_layers=args.num_layers,
mrl_k=args.mrl_k,
n_neurons=args.n_neurons,
window=args.window,
rho_star=args.rho_star,
surrogate=args.surrogate,
box_encoder=args.box_encoder,
ablation=args.ablation,
lstm_use_volume=args.lstm_use_volume,
lstm_hidden=args.lstm_hidden,
)
if args.structured_tokens:
# Replace the random frozen token table with the structured one so that
# low ids are genuinely "in-box" and high ids "out-of-box" -- this is
# what makes the concretization residual (channel 6) informative on the
# synthetic data. Stays frozen (requires_grad already False).
table = structured_token_table(args.vocab, args.d_box, args.n_in)
with torch.no_grad():
model.box_token_emb.weight.copy_(table)
return model
def build_optimizer(args, model):
params = [p for p in model.parameters() if p.requires_grad]
if args.optimizer == "adamw":
return torch.optim.AdamW(params, lr=args.lr, weight_decay=args.weight_decay)
if args.optimizer == "adam":
return torch.optim.Adam(params, lr=args.lr, weight_decay=args.weight_decay)
return torch.optim.SGD(params, lr=args.lr, momentum=0.9,
weight_decay=args.weight_decay)
def make_train_batch(args, seed: int):
"""A training batch: grounded (with optional overt/secret hallucinations)."""
return make_batch(
batch=args.batch_size, T=args.seq_len, d_model=args.d_model,
num_layers=args.num_layers, vocab=args.vocab, n_in=args.n_in,
topk=args.topk, hallu_prob=args.train_hallu_prob,
onset_lo=args.onset_lo, onset_hi=args.onset_hi,
lead_gap_lo=args.lead_gap_lo, lead_gap_hi=args.lead_gap_hi,
contaminate=args.contaminate, seed=seed, device=args.device,
)
def make_eval_batch(args, seed: int):
"""A held-out mixed batch with labels, for AUPRC + lead time."""
return make_batch(
batch=args.eval_batch, T=args.seq_len, d_model=args.d_model,
num_layers=args.num_layers, vocab=args.vocab, n_in=args.n_in,
topk=args.topk, hallu_prob=args.eval_hallu_prob,
onset_lo=args.onset_lo, onset_hi=args.onset_hi,
lead_gap_lo=args.lead_gap_lo, lead_gap_hi=args.lead_gap_hi,
contaminate=0.0, seed=seed, device=args.device,
)
@torch.no_grad()
def evaluate(args, model) -> tuple[dict, dict]:
"""Calibrate on grounded data, then score a mixed held-out batch.
Returns `(summary, detail)`: `summary` is the scalar metrics dict logged to
the console/checkpoints; `detail` exposes the per-step numpy arrays the
forward pass already computed (`anomaly`, `labels`, `burst`, `rho`, `t_star`,
`is_hallu`) so callers can build sample records without a second forward pass.
"""
model.eval()
# Calibrate the burst threshold on a purely grounded batch.
grounded, *_ = make_batch(
batch=args.eval_batch, T=args.seq_len, d_model=args.d_model,
num_layers=args.num_layers, vocab=args.vocab, n_in=args.n_in,
topk=args.topk, hallu_prob=0.0, contaminate=0.0,
seed=args.seed + 9991, device=args.device,
)
model.calibrate(grounded, z=args.calib_z)
signals, labels, t_star, is_hallu = make_eval_batch(args, seed=args.seed + 7777)
out = model(signals)
score = out.anomaly.detach().cpu().numpy()
y = labels.detach().cpu().numpy()
ap = auprc(score, y)
# Median lead time over the hallucinated sequences only.
#
# Primary metric uses `out.burst`, the *actual* detector decision
# (anomaly > anomaly_star AND rho > rho_star) -- so the reported warning
# corresponds to a token the detector would really fire on. We reuse the
# repo's `lead_time(series, t_star, threshold)` by passing the boolean burst
# as a 0/1 series thresholded at 0.5. We also report the rho-only lead time
# for reference, since that is what the bare `metrics.lead_time` measures.
burst = out.burst.float().detach().cpu().numpy()
rho = out.rho.detach().cpu().numpy()
ts = t_star.detach().cpu().numpy()
hal = is_hallu.detach().cpu().numpy().astype(bool)
rho_star = float(model.rho_star)
idx = [b for b in range(len(hal)) if hal[b]]
burst_leads = [lead_time(burst[b], int(ts[b]), 0.5) for b in idx]
rho_leads = [lead_time(rho[b], int(ts[b]), rho_star) for b in idx]
med_lead = float(np.median(burst_leads)) if burst_leads else float("nan")
med_lead_rho = float(np.median(rho_leads)) if rho_leads else float("nan")
summary = {
"auprc": float(ap),
"median_lead": med_lead, # from the calibrated burst detector
"median_lead_rho": med_lead_rho, # rho-only reference metric
"anomaly_star": float(model.anomaly_star),
"rho_star": rho_star,
"n_hallu": int(hal.sum()),
}
detail = {
"anomaly": score,
"labels": y,
"burst": burst,
"rho": rho,
"t_star": ts,
"is_hallu": hal,
}
return summary, detail
def _json_safe(x):
"""Recursively map NaN floats to None so json.dumps produces valid JSON.
metrics["auprc"]/median_lead* can be float("nan") (e.g. all-one-class eval
batches, or --eval-hallu-prob 0); plain json.dumps emits the invalid-JSON
literal `NaN` for those by default.
"""
if isinstance(x, float):
return None if x != x else x
if isinstance(x, dict):
return {k: _json_safe(v) for k, v in x.items()}
if isinstance(x, (list, tuple)):
return [_json_safe(v) for v in x]
return x
def _eval_samples_to_records(eval_detail: dict, n: int) -> list[dict]:
"""Map the first `n` eval-batch examples to scripts/visualize.py's record schema."""
anomaly = eval_detail["anomaly"]
labels = eval_detail["labels"]
burst = eval_detail["burst"]
rho = eval_detail["rho"]
t_star = eval_detail["t_star"]
is_hallu = eval_detail["is_hallu"]
B, T = anomaly.shape
n = min(n, B)
records = []
for b in range(n):
hallu = bool(is_hallu[b])
tstar_b = int(t_star[b])
question = (
f"synthetic eval example {b} "
f"({'hallucinated' if hallu else 'grounded'}"
f"{', t*=' + str(tstar_b) if hallu else ''})"
)
records.append({
"id": f"ex{b}",
"question": question,
"gen_tokens": [f"t{t}" for t in range(T)],
# The model's predicted per-step decision (not ground truth).
"token_labels": ["H" if burst[b, t] > 0.5 else "G" for t in range(T)],
"token_label_reasons": [
f"anomaly={anomaly[b, t]:.3f} rho={rho[b, t]:.3f}" for t in range(T)
],
# Ground-truth t*; synthetic.py uses T (out of range) as the "no t*"
# sentinel for grounded sequences, so this must be None, not T, or
# render_token_tree's tokens[t_star] indexing goes out of bounds.
"t_star_index": tstar_b if hallu else None,
})
return records
def _llm_epoch_eval(model, all_signals: list[dict], collate_fn, device: str,
batch_size: int, calib_z: float) -> tuple[dict, dict]:
"""Score every extracted LLM sample once against the auto-labelled ground
truth (token_labels / t_star -- real supervision, unlike the synthetic
pipeline's freshly generated held-out batch) and return (summary, detail),
mirroring evaluate()'s split.
anomaly_star has no meaning until calibrated (detector.py defaults it to
+inf). We calibrate on the mean/std of anomaly over every real
(non-padded) token across the *whole* dataset rather than a small tracked
subset -- a handful of samples lets a couple of outlier tokens blow up the
threshold so burst never fires again.
"""
was_training = model.training
model.eval()
with torch.no_grad():
flat_chunks = []
for start in range(0, len(all_signals), batch_size):
batch = all_signals[start:start + batch_size]
signals = collate_fn(batch, device)
anomaly = model(signals).anomaly.detach().cpu()
for i, s in enumerate(batch):
flat_chunks.append(anomaly[i, :s["meta"]["T"]])
flat = torch.cat(flat_chunks) if flat_chunks else torch.zeros(0)
model.anomaly_star = (flat.mean() + calib_z * flat.std()).detach().to(device)
scores, labels = [], []
kappa_grounded, kappa_hallucinated = [], []
burst_leads, rho_leads = [], []
per_sample: dict = {}
for start in range(0, len(all_signals), batch_size):
batch = all_signals[start:start + batch_size]
signals = collate_fn(batch, device)
out = model(signals)
anomaly = out.anomaly.detach().cpu().numpy()
rho = out.rho.detach().cpu().numpy()
kappa = out.channels.kappa.detach().cpu().numpy()
burst = out.burst.detach().cpu().numpy()
for i, s in enumerate(batch):
T = s["meta"]["T"]
true = s.get("token_labels") or []
# AUPRC over G/H tokens only; "U" (uncertain) is neither a
# clean positive nor negative, so it's excluded rather than
# silently counted as grounded.
for t in range(T):
lbl = true[t] if t < len(true) else None
if lbl == "U":
continue
scores.append(float(anomaly[i, t]))
labels.append(1 if lbl == "H" else 0)
if lbl == "H":
kappa_hallucinated.append(float(kappa[i, t]))
else:
kappa_grounded.append(float(kappa[i, t]))
t_star = s.get("t_star")
if t_star is not None and 0 <= t_star < T:
burst_leads.append(lead_time(burst[i, :T].astype(float), int(t_star), 0.5))
rho_leads.append(lead_time(rho[i, :T], int(t_star), float(model.rho_star)))
per_sample[s["record_id"]] = {
"anomaly": anomaly[i, :T], "rho": rho[i, :T], "burst": burst[i, :T],
}
if was_training:
model.train()
summary = {
"auprc": float(auprc(scores, labels)) if labels else float("nan"),
"median_lead": float(np.median(burst_leads)) if burst_leads else float("nan"),
"median_lead_rho": float(np.median(rho_leads)) if rho_leads else float("nan"),
"anomaly_star": float(model.anomaly_star),
"rho_star": float(model.rho_star),
"n_hallu_tokens": int(sum(labels)),
"mean_kappa_grounded": float(np.mean(kappa_grounded)) if kappa_grounded else float("nan"),
"mean_kappa_hallucinated": float(np.mean(kappa_hallucinated)) if kappa_hallucinated else float("nan"),
}
return summary, per_sample
def _llm_samples_to_records(viz_signals: list[dict], records_by_id: dict,
per_sample: dict) -> list[dict]:
"""Map a fixed LLM-sample subset's already-scored per-token predictions
(from _llm_epoch_eval's `detail`) into scripts/visualize.py's record
schema (see llm_signals.extract_all_signals for where record_id/t_star/
token_labels -- the auto-labelled ground truth -- come from)."""
records = []
for sample in viz_signals:
rid = sample["record_id"]
scored = per_sample.get(rid)
if scored is None:
continue
T = sample["meta"]["T"]
orig = records_by_id.get(rid, {})
true_labels = sample.get("token_labels") or []
true_t_star = sample.get("t_star")
anomaly, rho, burst = scored["anomaly"], scored["rho"], scored["burst"]
records.append({
"id": str(rid),
"question": orig.get("question", f"llm sample {rid}"),
"gen_tokens": [f"t{t}" for t in range(T)],
# The model's predicted per-step decision (not ground truth).
"token_labels": ["H" if burst[t] else "G" for t in range(T)],
"token_label_reasons": [
f"anomaly={anomaly[t]:.3f} rho={rho[t]:.3f} "
f"true={true_labels[t] if t < len(true_labels) else '?'}"
for t in range(T)
],
# Ground-truth t*; only meaningful if it indexes within this
# sample's actual token count (retokenized by --llm-model, which
# may not match the length of the original auto-labelled record).
"t_star_index": (true_t_star if true_t_star is not None and 0 <= true_t_star < T
else None),
})
return records
def append_history_line(out_dir: str, record: dict) -> None:
"""Append one JSON line to {out_dir}/history.jsonl (creates it on first write)."""
path = os.path.join(out_dir, "history.jsonl")
with open(path, "a") as f:
f.write(json.dumps(_json_safe(record)) + "\n")
def save_checkpoint(path: str, model, optimizer, args, epoch: int, metrics: dict):
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
torch.save({
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"args": vars(args),
"epoch": epoch,
"metrics": metrics,
}, path)
# ---------------------------------------------------------------------------
# training: synthetic generator
# ---------------------------------------------------------------------------
def train_synthetic(args, resume_ckpt: dict | None = None) -> dict:
"""Run training on the synthetic generator with already validated args."""
set_seed(args.seed)
device = resolve_device(args)
model = build_model(args).to(device)
optimizer = build_optimizer(args, model)
weights = LossWeights(vol=args.lam_vol, anchor=args.lam_anchor,
recon=args.lam_recon, mrl=args.lam_mrl,
join=args.lam_join, rank=args.lam_rank)
loss_fn = CombinedLoss(
model,
weights=weights,
trim=args.trim,
rank_margin=args.rank_margin,
rank_pair_mode=args.rank_pair_mode,
rank_max_pairs=None if args.rank_max_pairs == 0 else args.rank_max_pairs,
)
start_epoch = 0
best_ap = -1.0
if args.resume:
# weights_only=False: these are our own trusted local checkpoints and
# carry non-tensor metadata (args/metrics dicts). Explicit so the load is
# stable across torch versions that toggle the default.
if resume_ckpt is None:
resume_ckpt = torch.load(args.resume, map_location=device,
weights_only=False)
ckpt = resume_ckpt
model.load_state_dict(ckpt["model"])
if "optimizer" in ckpt:
optimizer.load_state_dict(ckpt["optimizer"])
start_epoch = int(ckpt.get("epoch", 0))
print(f"[resume] loaded {args.resume} (epoch {start_epoch}).")
# Carry the previous best AUPRC forward so resuming doesn't overwrite a
# better best.pt with a worse post-resume eval. Prefer the best.pt that
# sits in the output dir; fall back to the resumed checkpoint's metric.
best_path = os.path.join(args.out_dir, "best.pt")
prev_best = None
if os.path.exists(best_path):
prev_best = torch.load(best_path, map_location="cpu",
weights_only=False).get("metrics", {}).get("auprc")
if prev_best is None:
prev_best = ckpt.get("metrics", {}).get("auprc")
if prev_best is not None and prev_best == prev_best: # not NaN
best_ap = float(prev_best)
print(f"[resume] carrying best AUPRC={best_ap:.4f} forward.")
n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"[model] LatticeSpike d_box={args.d_box} n_neurons={args.n_neurons} "
f"mrl_k={args.mrl_k} surrogate={args.surrogate} | "
f"{n_params:,} trainable params | device={device}")
print(f"[data ] T={args.seq_len} d_model={args.d_model} L={args.num_layers} "
f"vocab={args.vocab} batch={args.batch_size} "
f"train_hallu={args.train_hallu_prob} contaminate={args.contaminate}")
if not args.no_save:
os.makedirs(args.out_dir, exist_ok=True)
with open(os.path.join(args.out_dir, "config.json"), "w") as f:
json.dump(vars(args), f, indent=2, sort_keys=True)
last_metrics: dict = {}
global_step = start_epoch * args.steps_per_epoch
t0 = time.time()
for epoch in range(start_epoch, start_epoch + args.epochs):
model.train()
running = {k: 0.0 for k in LOSS_KEYS}
for step in range(args.steps_per_epoch):
global_step += 1
signals, labels, *_ = make_train_batch(args, seed=args.seed + global_step)
out = model(signals)
loss, logs = loss_fn(out, signals["delta"], labels=labels)
optimizer.zero_grad(set_to_none=True)
loss.backward()
if args.grad_clip > 0:
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
optimizer.step()
for k in LOSS_KEYS:
running[k] += logs[k]
if args.log_every and (step + 1) % args.log_every == 0:
print(f" e{epoch:02d} s{step:03d} loss={logs['loss']:.4f} "
f"(contain={logs['contain']:.3f} vol={logs['vol']:.3f} "
f"anchor={logs['anchor']:.3f} recon={logs['recon']:.3f} "
f"mrl={logs['mrl']:.3f} join={logs['join']:.3f} "
f"rank={logs['rank']:.3f} "
f"gamma={logs['gamma']:.3f})")
avg_losses = {k: running[k] / max(1, args.steps_per_epoch) for k in LOSS_KEYS}
metrics, eval_detail = evaluate(args, model)
last_metrics = metrics
elapsed = time.time() - t0
print(f"[epoch {epoch:02d}] train_loss={avg_losses['loss']:.4f} "
f"AUPRC={metrics['auprc']:.4f} "
f"lead(burst)={metrics['median_lead']:.2f} "
f"lead(rho)={metrics['median_lead_rho']:.2f} "
f"anomaly*={metrics['anomaly_star']:.4f} "
f"({elapsed:.1f}s)")
# Track best AUPRC regardless of whether we persist checkpoints
# (NaN-safe: AUPRC is NaN only when a batch has no positives).
is_best = metrics["auprc"] == metrics["auprc"] and metrics["auprc"] > best_ap
if is_best:
best_ap = metrics["auprc"]
if not args.no_save:
epoch_num = epoch + 1
history_record = {
"epoch": epoch_num,
"elapsed_sec": elapsed,
"train_loss": avg_losses,
"eval": metrics,
}
if args.viz_samples > 0 and epoch_num % args.viz_every == 0:
history_record["samples"] = _eval_samples_to_records(
eval_detail, args.viz_samples)
append_history_line(args.out_dir, history_record)
save_checkpoint(os.path.join(args.out_dir, "last.pt"),
model, optimizer, args, epoch_num, metrics)
if is_best:
save_checkpoint(os.path.join(args.out_dir, "best.pt"),
model, optimizer, args, epoch_num, metrics)
if args.save_every and epoch_num % args.save_every == 0:
save_checkpoint(
os.path.join(args.out_dir, f"epoch{epoch_num:03d}.pt"),
model, optimizer, args, epoch_num, metrics)
print(f"[done] best AUPRC={best_ap:.4f} "
f"checkpoints in {args.out_dir!r}" if not args.no_save
else f"[done] best AUPRC={best_ap:.4f} (no checkpoints written)")
return last_metrics
# ---------------------------------------------------------------------------
# training: real LLM signals
# ---------------------------------------------------------------------------
def _pad_to(tensor_list: list[torch.Tensor], T_max: int) -> torch.Tensor:
"""Pad each (B, T, ...) tensor to T_max along dim 1 and cat on dim 0."""
padded = []
for t in tensor_list:
pad_len = T_max - t.shape[1]
if pad_len > 0:
pad_shape = list(t.shape)
pad_shape[1] = pad_len
t = torch.cat([t, torch.zeros(pad_shape, dtype=t.dtype)], dim=1)
padded.append(t)
return torch.cat(padded, dim=0)
def train_llm(args) -> dict:
"""Extract real LLM signals from a JSONL dataset and train LatticeSpike.
Unlike the synthetic pipeline (fresh generated batch per step, calibrated
against a separately generated held-out batch), this trains over shuffled
minibatches of a fixed, pre-extracted dataset and evaluates AUPRC/lead-
time each epoch against the dataset's own auto-labelled token_labels/
t_star (see _llm_epoch_eval), calibrating anomaly_star over the whole set
rather than a single batch. Per-epoch losses, eval metrics, and (on
tracked epochs) decision-tree sample snapshots (_llm_samples_to_records)
are appended to {out-dir}/history.jsonl for scripts/visualize_train.py,
then the final RSNNOutput is dumped for offline inspection instead of
maintaining last.pt/best.pt checkpoints.
"""
if args.lam_rank > 0.0:
raise ValueError(
"supervised ranking is not available in train.py --data-source llm "
"because this legacy path cannot prove token-label alignment; use "
"scripts/train_latticespike_llm_eval.py instead"
)
set_seed(args.seed)
device = resolve_device(args)
print(f"[config] data-source=llm device={device} epochs={args.epochs} "
f"batch_size={args.batch_size}")
# requires `transformers`; imported lazily so the synthetic pipeline (and
# the rest of the package) runs without it installed.
from transformers import AutoModelForCausalLM, AutoTokenizer
llm_signals = __import__(
f"{_PKG}.llm_signals",
fromlist=["load_jsonl", "extract_all_signals", "collate_batch"],
)
load_jsonl = llm_signals.load_jsonl
extract_all_signals = llm_signals.extract_all_signals
collate_batch = llm_signals.collate_batch
def resolve_path(p: str) -> str:
return p if os.path.isabs(p) else os.path.join(_HERE, p)
data_path = resolve_path(args.llm_data)
out_dir = resolve_path(args.out_dir)
os.makedirs(out_dir, exist_ok=True)
print(f"[load ] LLM: {args.llm_model}")
tokenizer = AutoTokenizer.from_pretrained(args.llm_model)
llm = AutoModelForCausalLM.from_pretrained(
args.llm_model, dtype=torch.float32, attn_implementation="eager"
)
llm = llm.to(device).eval()
llm_cfg = llm.config
d_model: int = llm_cfg.hidden_size
num_layers: int = llm_cfg.num_hidden_layers
vocab_size: int = llm_cfg.vocab_size
print(f"[data ] loading {data_path}")
records = load_jsonl(data_path)
records_by_id = {rec.get("id", i): rec for i, rec in enumerate(records)}
print(f"[data ] {len(records)} records")
print("[extract] running LLM forward passes (this may take a few minutes)...")
t0 = time.time()
all_signals = extract_all_signals(
records, llm, tokenizer, device,
topk=args.topk, text_field=args.text_field, limit=args.limit,
)
print(f"[extract] {len(all_signals)} samples in {time.time() - t0:.1f}s")
viz_signals = all_signals[:args.viz_samples] if args.viz_samples > 0 else []
# Free LLM (GPU) memory before the training model is allocated.
del llm
if device.startswith("cuda"):
torch.cuda.empty_cache()
print(f"[model ] d_model={d_model} num_layers={num_layers} vocab_size={vocab_size}")
print(f"[model ] d_box={args.d_box} n_neurons={args.n_neurons} mrl_k={args.mrl_k}")
# Real dims come from the loaded LM, not the synthetic-generator flags.
args.d_model, args.num_layers, args.vocab = d_model, num_layers, vocab_size
args.structured_tokens = False
model = build_model(args).to(device)
n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"[model ] {n_params:,} trainable params")
optimizer = build_optimizer(args, model)
weights = LossWeights(vol=args.lam_vol, anchor=args.lam_anchor,
recon=args.lam_recon, mrl=args.lam_mrl,
join=args.lam_join, rank=args.lam_rank)
loss_fn = CombinedLoss(
model,
weights=weights,
trim=args.trim,
rank_margin=args.rank_margin,
rank_pair_mode=args.rank_pair_mode,
rank_max_pairs=None if args.rank_max_pairs == 0 else args.rank_max_pairs,
)
n = len(all_signals)
print(f"[train ] {args.epochs} epochs | {n} samples | batch_size={args.batch_size}")
t0 = time.time()
avg_losses = {k: float("nan") for k in LOSS_KEYS}
for epoch in range(args.epochs):
model.train()
idx = torch.randperm(n).tolist()
running = {k: 0.0 for k in LOSS_KEYS}
steps = 0
for start in range(0, n, args.batch_size):
batch = [all_signals[i] for i in idx[start:start + args.batch_size]]
signals = collate_batch(batch, device)
out = model(signals)
loss, logs = loss_fn(out, signals["delta"])
optimizer.zero_grad(set_to_none=True)
loss.backward()
if args.grad_clip > 0:
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
optimizer.step()
for k in LOSS_KEYS:
running[k] += logs[k]
steps += 1
avg_losses = {k: running[k] / max(1, steps) for k in LOSS_KEYS}
elapsed = time.time() - t0
eval_summary, per_sample = _llm_epoch_eval(
model, all_signals, collate_batch, device, args.batch_size, args.calib_z)
print(f" epoch {epoch:02d} avg_loss={avg_losses['loss']:.4f} "
f"AUPRC={eval_summary['auprc']:.4f} "
f"lead(burst)={eval_summary['median_lead']:.2f} "
f"lead(rho)={eval_summary['median_lead_rho']:.2f} "
f"({elapsed:.1f}s elapsed)")
epoch_num = epoch + 1
history_record = {
"epoch": epoch_num,
"elapsed_sec": elapsed,
"train_loss": avg_losses,
"eval": eval_summary,
}
if viz_signals and epoch_num % args.viz_every == 0:
history_record["samples"] = _llm_samples_to_records(viz_signals, records_by_id, per_sample)
append_history_line(out_dir, history_record)
# Final eval forward pass -- collect RSNNOutput over all samples.
print("[eval ] running final forward pass on all samples...")
model.eval()
all_spikes, all_u, all_u_pred, all_gates = [], [], [], []
all_rho, all_anomaly, all_burst = [], [], []
all_c_min, all_c_max = [], []
all_kappa, all_r = [], []
seq_lengths: list[int] = []
with torch.no_grad():
for start in range(0, n, args.batch_size):
batch = all_signals[start:start + args.batch_size]
signals = collate_batch(batch, device)
out = model(signals)
r = out.rsnn
all_spikes.append(r.spikes.cpu())
all_u.append(r.u.cpu())
all_u_pred.append(r.u_pred.cpu())
all_gates.append(r.gates.cpu())
all_rho.append(out.rho.cpu())
all_anomaly.append(out.anomaly.cpu())
all_burst.append(out.burst.cpu())
all_c_min.append(out.box.c_min.cpu())
all_c_max.append(out.box.c_max.cpu())
all_kappa.append(out.channels.kappa.cpu())
all_r.append(out.channels.leakage.cpu())
seq_lengths.extend([s["meta"]["T"] for s in batch])
# T may vary across mini-batches after padding, so pad to a global T_max.
T_max = max(s["meta"]["T"] for s in all_signals)
spikes_all = _pad_to(all_spikes, T_max)
u_all = _pad_to(all_u, T_max)
u_pred_all = _pad_to(all_u_pred, T_max)
gates_all = _pad_to(all_gates, T_max)
rho_all = _pad_to(all_rho, T_max)
anomaly_all = _pad_to(all_anomaly, T_max)
burst_all = _pad_to(all_burst, T_max)
c_min_all = _pad_to(all_c_min, T_max)
c_max_all = _pad_to(all_c_max, T_max)
kappa_all = _pad_to(all_kappa, T_max)
r_all = _pad_to(all_r, T_max)
out_path = os.path.join(out_dir, "rsnn_output.pt")
torch.save(
{
"spikes": spikes_all, "u": u_all, "u_pred": u_pred_all, "gates": gates_all,
"rho": rho_all, "anomaly": anomaly_all, "burst": burst_all,
"c_min": c_min_all, "c_max": c_max_all,
"kappa": kappa_all, "r": r_all,
"seq_lengths": seq_lengths,
"record_ids": [s["record_id"] for s in all_signals],
"t_star": [s["t_star"] for s in all_signals],
"token_labels": [s["token_labels"] for s in all_signals],
"config": {
"d_model": d_model, "vocab_size": vocab_size, "num_layers": num_layers,
"d_box": args.d_box, "n_neurons": args.n_neurons, "mrl_k": args.mrl_k,
"topk": args.topk, "epochs": args.epochs,
"box_encoder": args.box_encoder,
"lam_rank": args.lam_rank,
"rank_margin": args.rank_margin,
"rank_pair_mode": args.rank_pair_mode,
"rank_max_pairs": args.rank_max_pairs,
"ablation": args.ablation, "lstm_use_volume": args.lstm_use_volume,
"lstm_hidden": args.lstm_hidden,
},
},
out_path,
)
ckpt_path = os.path.join(out_dir, "rsnn_checkpoint.pt")
save_checkpoint(ckpt_path, model, optimizer, args, args.epochs, {"avg_loss": avg_losses["loss"]})
print(f"[save ] RSNNOutput -> {out_path}")
print(f" spikes : {tuple(spikes_all.shape)}")
print(f" u : {tuple(u_all.shape)}")
print(f" u_pred : {tuple(u_pred_all.shape)}")
print(f" gates : {tuple(gates_all.shape)}")
print(f" c_min : {tuple(c_min_all.shape)}")
print(f" c_max : {tuple(c_max_all.shape)}")
print(f" kappa : {tuple(kappa_all.shape)}")
print(f" r : {tuple(r_all.shape)}")
print(f"[save ] checkpoint -> {ckpt_path}")
print("[done ]")
return {"avg_loss": avg_losses["loss"]}
def main(argv=None) -> None:
parser = build_parser()
args = parser.parse_args(argv)
overrides = cli_overrides(parser, argv)
if args.data_source == "llm":
apply_defaults(args, LLM_MODE_DEFAULTS, overrides)
errs = validate_args(args)
if errs:
parser.error("invalid configuration:\n - " + "\n - ".join(errs))
train_llm(args)
return
resume_ckpt = None
if args.resume:
try:
resume_ckpt = torch.load(args.resume, map_location="cpu",
weights_only=False)
except Exception as exc:
parser.error(f"could not load --resume checkpoint {args.resume!r}: {exc}")
apply_checkpoint_args(args, resume_ckpt, overrides)
errs = validate_args(args)
if errs:
parser.error("invalid configuration:\n - " + "\n - ".join(errs))
train_synthetic(args, resume_ckpt=resume_ckpt)
if __name__ == "__main__":
main()
"""
LatticeSpike training driver.
Restores the `train_demo.py` entry point referenced by the README (see
PROJECT_IDEA.md, gap #1): it wires a data source into the end-to-end
LatticeSpike model and trains the combined unsupervised objective.
``--data-source`` selects which pipeline feeds it:
synthetic (default) : the synthetic signal generator (synthetic.py).
Calibrates the anomaly threshold and evaluates
AUPRC + lead time each epoch.
llm : real signals extracted from a HuggingFace causal
LM (llm_signals.py) over a labelled JSONL dataset.
Trains over the whole dataset each epoch and dumps
the final RSNNOutput for offline inspection.
Every internal setting of the run is exposed as a command-line flag (model
architecture, data generation/extraction, loss weights, optimization
schedule, calibration, and checkpointing), so a run can be reconfigured
without touching code.
This module is imported as part of the ``Lattice_Spike`` package (its siblings
use ``from .boxes import ...``), so the block below puts the package's parent
directory on ``sys.path`` and then imports the package by name. That makes
python Lattice_Spike/train.py # from the repo's parent dir
python train.py # from inside the package dir
python -m Lattice_Spike.train # as a module
all work the same way.
Examples
--------
# Quick CPU demo with defaults (synthetic data).
python Lattice_Spike/train.py
# Paper-scale boxes, sigmoid surrogate, robust to 15% contamination.
python Lattice_Spike/train.py --d-box 256 --n-neurons 512 \
--surrogate sigmoid --contaminate 0.15 --trim 0.2
# Resume from a checkpoint and train 5 more epochs.
python Lattice_Spike/train.py --resume runs/demo/last.pt --epochs 5
# Train on real LLM signals extracted from a labelled JSONL dataset.
python Lattice_Spike/train.py --data-source llm --limit 100 --epochs 3
"""
from __future__ import annotations
import argparse
import json
import os
import sys
import time
# --- make the package importable however this file is launched --------------
_HERE = os.path.dirname(os.path.abspath(__file__))
_PARENT = os.path.dirname(_HERE)
_PKG = os.path.basename(_HERE) # the package directory name
if _PARENT not in sys.path:
sys.path.insert(0, _PARENT)
import numpy as np
import torch
_pkg = __import__(_PKG) # e.g. import Lattice_Spike
LatticeSpike = _pkg.LatticeSpike
CombinedLoss = _pkg.CombinedLoss
LossWeights = _pkg.LossWeights
auprc = _pkg.auprc
lead_time = _pkg.lead_time
make_batch = __import__(f"{_PKG}.synthetic", fromlist=["make_batch"]).make_batch
structured_token_table = __import__(
f"{_PKG}.synthetic", fromlist=["structured_token_table"]
).structured_token_table
# llm_signals is imported lazily inside train_llm() -- synthetic-only runs
# (the common case) shouldn't pay for loading the LLM-signal-extraction
# pipeline at all.
# Defaults that differ from the synthetic-mode flag defaults below, applied to
# --data-source llm runs unless the user explicitly overrides the flag.
LLM_MODE_DEFAULTS = {
"topk": 64,
"batch_size": 16,
"epochs": 5,
"lr": 1e-3,
"n_neurons": 64,
"box_encoder": "linear",
"out_dir": "output",
}
# Stable loss schema retained so anchor-off runs remain directly comparable
# with historical runs.
LOSS_KEYS = ("loss", "contain", "vol", "anchor", "recon", "mrl", "gamma")
# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------