| """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()} |
|
|
