File size: 4,523 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 | """General decision datasets in JSONL, one decision per line:
{"state": "hotdog", "question": "Is this food?", "choices": ["yes", "no"], "target": "yes"}
Optional fields:
teacher_distribution {"yes": 0.97, "no": 0.03} soft labels from a teacher (LLM, oracle, ...)
teacher_probs [0.97, 0.03] same, aligned with `choices`
score 0.83 continuous target for the score head
group any id; splits and dedupe work per group (entity, document, episode)
`target` may be the choice text or its index; it may be missing when only a
teacher distribution is given. The teacher is only needed to *build* the data:
distillation stores its probabilities, the deployed model never calls it.
"""
from __future__ import annotations
import json
import zlib
from pathlib import Path
from typing import Iterable
import numpy as np
from .decision_data import DecisionData
NEG = -1e9
def read_jsonl(path: str | Path) -> list[dict]:
with open(path, encoding="utf-8") as f:
return [json.loads(line) for line in f if line.strip()]
def write_jsonl(path: str | Path, rows: Iterable[dict]) -> None:
Path(path).parent.mkdir(parents=True, exist_ok=True)
with open(path, "w", encoding="utf-8") as f:
for r in rows:
f.write(json.dumps(r, ensure_ascii=False) + "\n")
def _gid(v) -> int:
if isinstance(v, (int, np.integer)):
return int(v)
return zlib.crc32(str(v).encode())
def teacher_probs(row: dict) -> list[float] | None:
if row.get("teacher_probs") is not None:
p = [float(x) for x in row["teacher_probs"]]
elif row.get("teacher_distribution") is not None:
td = row["teacher_distribution"]
p = [float(td.get(c, 0.0)) for c in row["choices"]]
else:
return None
s = sum(p)
if len(p) != len(row["choices"]) or s <= 0:
raise ValueError(f"bad teacher distribution for row {row.get('question')!r}")
return [x / s for x in p]
def rows_to_data(rows: list[dict], task: str = "text", task_cfg: dict | None = None) -> DecisionData:
from minijev.tasks import get_task
t = get_task(task, task_cfg or {})
n = len(rows)
nmax = max((len(r.get("choices") or []) for r in rows), default=1) or 1
qvocab, cvocab = {}, {}
qi = np.zeros(n, np.int32)
ci = np.full((n, nmax), -1, np.int32)
target = np.full(n, -1, np.int64)
teacher = np.full((n, nmax), NEG, np.float32)
score = np.full(n, np.nan, np.float32)
has_teacher = has_score = False
for i, r in enumerate(rows):
qi[i] = qvocab.setdefault(r.get("question", ""), len(qvocab))
ch = list(r.get("choices") or [])
if len(set(ch)) != len(ch):
raise ValueError(f"duplicate choices in row {i}: {ch}")
for j, c in enumerate(ch):
ci[i, j] = cvocab.setdefault(c, len(cvocab))
tg = r.get("target")
if tg is not None:
target[i] = tg if isinstance(tg, int) else ch.index(tg)
p = teacher_probs(r)
if p is not None:
has_teacher = True
teacher[i, : len(p)] = np.log(np.maximum(np.array(p, np.float32), 1e-8))
if r.get("score") is not None:
has_score = True
score[i] = float(r["score"])
groups = np.array([_gid(r.get("group", r.get("state") or r.get("question", ""))) for r in rows], np.int64)
hashes = np.array([_gid(f"{r.get('state') or ''}\x00{r.get('question', '')}") for r in rows], np.int64)
return DecisionData(task=task, state=t.states_to_arrays([r.get("state") or "" for r in rows]),
questions=list(qvocab), question_idx=qi, choices=list(cvocab), choice_idx=ci,
target=target, group=groups, teacher_logits=teacher if has_teacher else None,
hash=hashes, score=score if has_score else None,
meta={"source": "jsonl", "rows": n})
def generate_splits(out_dir: str | Path, samples: int = 0, source: str | None = None, seed: int = 0,
fractions=(0.9, 0.05, 0.05), task: str = "text", **_) -> dict:
"""`data.generate: {source: file.jsonl}` -> group-split train/val/test .npz."""
from .decision_data import split_by_group
data = rows_to_data(read_jsonl(source), task=task)
splits = split_by_group(data, fractions, seed=seed)
for k, v in splits.items():
v.save(Path(out_dir) / f"{k}.npz")
return {k: len(v) for k, v in splits.items()}
|