"""Generic on-disk dataset of decisions, independent of the task. Each row: a task-specific state (stored by the task plugin as numpy arrays, ragged arrays as `X` + `X_offsets`), a question, a variable-size list of choices and a supervision signal: target hard label (index into the row's choices), -1 if none soft optional soft labels (n, Nmax) -> soft-label training teacher_logits optional teacher logits (n, Nmax) -> knowledge distillation group episode / document id: splits are made by group (no leakage) hash optional exact-state hash to drop duplicates across splits score optional continuous target in [0, 1] (NaN = none) for `model.score()` """ from __future__ import annotations import json from dataclasses import dataclass, field from pathlib import Path import numpy as np from minijev.tasks import ragged_take OPTIONAL = ("soft", "teacher_logits", "hash", "score") @dataclass class DecisionData: task: str state: dict[str, np.ndarray] questions: list[str] question_idx: np.ndarray choices: list[str] choice_idx: np.ndarray # (n, Nmax) int32, -1 = no choice target: np.ndarray group: np.ndarray soft: np.ndarray | None = None teacher_logits: np.ndarray | None = None hash: np.ndarray | None = None score: np.ndarray | None = None # (n,) float in [0, 1], NaN = none -> trains the scalar score head meta: dict = field(default_factory=dict) def __len__(self) -> int: return len(self.target) # ------------------------------------------------------------------ io def save(self, path: str | Path) -> None: path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) arrays = {f"state/{k}": v for k, v in self.state.items()} arrays.update(question_idx=self.question_idx, choice_idx=self.choice_idx, target=self.target, group=self.group) for k in OPTIONAL: if getattr(self, k) is not None: arrays[k] = getattr(self, k) header = {"task": self.task, "questions": self.questions, "choices": self.choices, "meta": self.meta} arrays["header"] = np.frombuffer(json.dumps(header).encode(), dtype=np.uint8) np.savez(path, **arrays) @classmethod def load(cls, path: str | Path) -> "DecisionData": z = np.load(path) header = json.loads(bytes(z["header"]).decode()) state = {k[6:]: z[k] for k in z.files if k.startswith("state/")} opt = {k: z[k] for k in OPTIONAL if k in z.files} return cls(header["task"], state, header["questions"], z["question_idx"], header["choices"], z["choice_idx"], z["target"], z["group"], meta=header.get("meta", {}), **opt) # ------------------------------------------------------------ slicing def subset(self, idx: np.ndarray) -> "DecisionData": idx = np.asarray(idx, dtype=np.int64) state = {} for k, v in self.state.items(): if k.endswith("_offsets"): continue if f"{k}_offsets" in self.state: flat, lengths = ragged_take(v, self.state[f"{k}_offsets"], idx) offs = np.zeros(len(idx) + 1, dtype=np.int64) np.cumsum(lengths, out=offs[1:]) state[k], state[f"{k}_offsets"] = flat, offs else: state[k] = v[idx] opt = {k: getattr(self, k)[idx] for k in OPTIONAL if getattr(self, k) is not None} return DecisionData(self.task, state, self.questions, self.question_idx[idx], self.choices, self.choice_idx[idx], self.target[idx], self.group[idx], meta=dict(self.meta), **opt) @classmethod def concat(cls, parts: list["DecisionData"]) -> "DecisionData": """Rows of several datasets of the same task; question/choice vocabularies are merged.""" first = parts[0] questions = list(dict.fromkeys(q for p in parts for q in p.questions)) choices = list(dict.fromkeys(c for p in parts for c in p.choices)) qmap = {q: i for i, q in enumerate(questions)} cmap = {c: i for i, c in enumerate(choices)} nmax = max(p.choice_idx.shape[1] for p in parts) qi, ci = [], [] for p in parts: qi.append(np.array([qmap[q] for q in p.questions], np.int32)[p.question_idx]) clut = np.array([cmap[c] for c in p.choices] + [-1], np.int32) # index -1 -> -1 pad = np.full((len(p), nmax), -1, np.int32) pad[:, : p.choice_idx.shape[1]] = clut[p.choice_idx] ci.append(pad) state = {} for k in first.state: if k.endswith("_offsets"): continue if f"{k}_offsets" in first.state: offs, base = [np.zeros(1, np.int64)], 0 for p in parts: offs.append(p.state[f"{k}_offsets"][1:] + base) base += int(p.state[f"{k}_offsets"][-1]) state[f"{k}_offsets"] = np.concatenate(offs) state[k] = np.concatenate([p.state[k] for p in parts]) opt = {} for k in OPTIONAL: if getattr(first, k) is None: continue vals = [] for p in parts: v = getattr(p, k) if v.ndim == 2 and v.shape[1] < nmax: fill = -1e9 if k == "teacher_logits" else 0.0 v = np.pad(v, ((0, 0), (0, nmax - v.shape[1])), constant_values=fill) vals.append(v) opt[k] = np.concatenate(vals) return cls(first.task, state, questions, np.concatenate(qi), choices, np.concatenate(ci), np.concatenate([p.target for p in parts]), np.concatenate([p.group for p in parts]), meta=dict(first.meta), **opt) def row_choices(self, i: int) -> list[str]: return [self.choices[j] for j in self.choice_idx[i] if j >= 0] def split_by_group(data: DecisionData, fractions=(0.9, 0.05, 0.05), seed: int = 0, dedupe: bool = True) -> dict[str, DecisionData]: """Group-level split (whole episodes/documents go to one split), then drop val/test rows whose exact state also appears in train.""" groups = np.unique(data.group) rng = np.random.default_rng(seed) rng.shuffle(groups) n_val = max(1, int(round(len(groups) * fractions[1]))) n_test = max(1, int(round(len(groups) * fractions[2]))) sets = {"val": groups[:n_val], "test": groups[n_val:n_val + n_test], "train": groups[n_val + n_test:]} out, dropped = {}, {} for name in ("train", "val", "test"): idx = np.nonzero(np.isin(data.group, sets[name]))[0] if dedupe and name != "train" and data.hash is not None: keep = ~np.isin(data.hash[idx], data.hash[out["train_idx"]]) dropped[name] = int((~keep).sum()) idx = idx[keep] out[f"{name}_idx"] = idx splits = {name: data.subset(out[f"{name}_idx"]) for name in ("train", "val", "test")} for name, s in splits.items(): s.meta.update(split=name, dropped_duplicates=dropped.get(name, 0)) return splits