"""On-disk feature cache: one safetensors file per example plus a JSONL manifest.""" from __future__ import annotations import json from dataclasses import asdict, dataclass, field from pathlib import Path from typing import Optional import torch from safetensors.torch import load_file, save_file from .schema import TYPE_INDEX @dataclass class Decision: qid: str type: str slot: int spans: list[list[int]] alias: list[int] prior: list[float] gold: Optional[list[float]] @dataclass class Record: id: str source: str path: str n_prompt: int n_canvas: int decisions: list[Decision] = field(default_factory=list) def to_json(self) -> str: return json.dumps(asdict(self)) @classmethod def from_json(cls, line: str) -> "Record": d = json.loads(line) d["decisions"] = [Decision(**x) for x in d["decisions"]] return cls(**d) class CacheWriter: """Writes one safetensors file per example and a JSONL manifest. A writer always starts a fresh manifest, so re-running into the same folder does not duplicate records. """ def __init__(self, out_dir): self.dir = Path(out_dir) self.dir.mkdir(parents=True, exist_ok=True) self._f = (self.dir / "manifest.jsonl").open("w", encoding="utf-8") def write(self, record: Record, feats: torch.Tensor) -> None: save_file({"feats": feats.detach().to("cpu", torch.float16).contiguous()}, str(self.dir / record.path)) self._f.write(record.to_json() + "\n") self._f.flush() def close(self) -> None: self._f.close() class FeatureCache: def __init__(self, cache_dir): self.dir = Path(cache_dir) with (self.dir / "manifest.jsonl").open("r", encoding="utf-8") as f: self.records = [Record.from_json(line) for line in f if line.strip()] def __len__(self) -> int: return len(self.records) def __getitem__(self, i: int) -> tuple[Record, torch.Tensor]: rec = self.records[i] return rec, load_file(str(self.dir / rec.path))["feats"] @dataclass class Batch: """A padded feature batch plus two views of its decisions. `decisions` is the per-decision list (metadata, spans, prior, gold) that evaluation, serving and the tests read. The flat tensors are what the head and the losses compute on: options are numbered decision by decision in that same order, so one gather, one index_add and one padded log-softmax cover the whole batch instead of a Python loop issuing hundreds of tiny ops per decision (session 3's CPU-bound trainer). """ feats: torch.Tensor seg: torch.Tensor mask: torch.Tensor decisions: list[dict] dec_b: torch.Tensor # (n_dec,) batch row of each decision dec_slot: torch.Tensor # (n_dec,) answer-slot row of each decision dec_type: torch.Tensor # (n_dec,) question type id dec_has_gold: torch.Tensor # (n_dec,) bool opt_dec: torch.Tensor # (n_opt,) decision index of each option opt_pos: torch.Tensor # (n_opt,) position of each option within its decision opt_prior: torch.Tensor # (n_opt,) backbone log-prior per option opt_gold: torch.Tensor # (n_opt,) gold probability per option, 0 where the decision has no gold span_rows: torch.Tensor # (n_span_rows,) row index into the padded (B * L) feature rows, option by option span_opt: torch.Tensor # (n_span_rows,) option index of each span row span_len: torch.Tensor # (n_opt,) span length per option, float counts: list[int] # options per decision; plain Python so torch.split never syncs on it n_gold: int # decisions with gold; plain Python so the losses never sync on it _FLAT = ("dec_b", "dec_slot", "dec_type", "dec_has_gold", "opt_dec", "opt_pos", "opt_prior", "opt_gold", "span_rows", "span_opt", "span_len") def to(self, device) -> "Batch": # spans stay on the CPU: consumers read them back as Python ints, so a device copy would only add a # device-to-host sync per decision. return Batch(self.feats.to(device), self.seg.to(device), self.mask.to(device), [{**d, "prior": d["prior"].to(device), "gold": None if d["gold"] is None else d["gold"].to(device)} for d in self.decisions], *(getattr(self, name).to(device) for name in self._FLAT), self.counts, self.n_gold) def _check_record(rec: Record, n_rows: int) -> None: """Raise a ValueError naming the record and question when a record does not fit its feature rows. Rows [0, n_prompt) hold the prompt and [n_prompt, n_prompt + n_canvas) the canvas. Without this check a bad record would be read silently from padding or from the wrong segment. """ if n_rows != rec.n_prompt + rec.n_canvas: raise ValueError(f"record {rec.id!r}: {n_rows} feature rows, but n_prompt {rec.n_prompt} + n_canvas " f"{rec.n_canvas} = {rec.n_prompt + rec.n_canvas}") for d in rec.decisions: where = f"record {rec.id!r}, question {d.qid!r}" if not rec.n_prompt <= d.slot < n_rows: raise ValueError(f"{where}: slot {d.slot} is outside the canvas rows [{rec.n_prompt}, {n_rows})") for s, e in d.spans: if not 0 <= s < e <= rec.n_prompt: raise ValueError(f"{where}: option span [{s}, {e}) is empty or outside the prompt rows " f"[0, {rec.n_prompt})") if len(d.prior) != len(d.spans): raise ValueError(f"{where}: {len(d.spans)} option spans but {len(d.prior)} prior values") if d.gold is not None and len(d.gold) != len(d.spans): raise ValueError(f"{where}: {len(d.spans)} option spans but {len(d.gold)} gold values") def collate(items: list[tuple[Record, torch.Tensor]]) -> Batch: lmax = max(f.shape[0] for _, f in items) dim = items[0][1].shape[1] feats = torch.zeros(len(items), lmax, dim, dtype=torch.float16) seg = torch.zeros(len(items), lmax, dtype=torch.long) mask = torch.zeros(len(items), lmax, dtype=torch.bool) decisions = [] dec_b, dec_slot, dec_type, dec_has_gold, counts = [], [], [], [], [] opt_dec, opt_pos, opt_prior, opt_gold, span_rows, span_opt, span_len = [], [], [], [], [], [], [] for b, (rec, f) in enumerate(items): n = f.shape[0] _check_record(rec, n) feats[b, :n] = f.to(torch.float16) mask[b, :n] = True seg[b, rec.n_prompt:n] = 1 for d in rec.decisions: i = len(decisions) decisions.append({ "b": b, "slot": d.slot, "type_id": TYPE_INDEX[d.type], "qid": d.qid, "type": d.type, "spans": torch.tensor(d.spans, dtype=torch.long), "prior": torch.tensor(d.prior, dtype=torch.float32), "gold": None if d.gold is None else torch.tensor(d.gold, dtype=torch.float32), }) dec_b.append(b) dec_slot.append(d.slot) dec_type.append(TYPE_INDEX[d.type]) dec_has_gold.append(d.gold is not None) counts.append(len(d.spans)) for j, (s, e) in enumerate(d.spans): o = len(opt_dec) opt_dec.append(i) opt_pos.append(j) opt_prior.append(d.prior[j]) opt_gold.append(0.0 if d.gold is None else d.gold[j]) span_rows.extend(range(b * lmax + s, b * lmax + e)) span_opt.extend([o] * (e - s)) span_len.append(float(e - s)) long, f32 = torch.long, torch.float32 return Batch(feats=feats, seg=seg, mask=mask, decisions=decisions, dec_b=torch.tensor(dec_b, dtype=long), dec_slot=torch.tensor(dec_slot, dtype=long), dec_type=torch.tensor(dec_type, dtype=long), dec_has_gold=torch.tensor(dec_has_gold, dtype=torch.bool), opt_dec=torch.tensor(opt_dec, dtype=long), opt_pos=torch.tensor(opt_pos, dtype=long), opt_prior=torch.tensor(opt_prior, dtype=f32), opt_gold=torch.tensor(opt_gold, dtype=f32), span_rows=torch.tensor(span_rows, dtype=long), span_opt=torch.tensor(span_opt, dtype=long), span_len=torch.tensor(span_len, dtype=f32), counts=counts, n_gold=sum(dec_has_gold))