| """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 |
| 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 |
| meta: dict = field(default_factory=dict) |
|
|
| def __len__(self) -> int: |
| return len(self.target) |
|
|
| |
| 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) |
|
|
| |
| 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) |
| 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 |
|
|
