File size: 7,229 Bytes
026f63e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 | """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
|