yashu2000/TemporalBenchEnv

1"""Load question banks from JSON or JSONL files."""2 3from __future__ import annotations4 5import json6from pathlib import Path7from typing import Any8 9from .question import TSQuestion10 11# Canonical domain keys used by EpisodeSampler (must match bank files or dataset field)12DEFAULT_DOMAIN_ORDER = ("PSML", "freshretailnet", "MIMIC", "causal_chambers")13 14 15def _parse_records(raw: Any) -> list[dict[str, Any]]:16    if isinstance(raw, list):17        return [x for x in raw if isinstance(x, dict)]18    if isinstance(raw, dict) and "questions" in raw:19        q = raw["questions"]20        if isinstance(q, list):21            return [x for x in q if isinstance(x, dict)]22    raise ValueError("JSON root must be a list of objects or {\"questions\": [...]}")23 24 25def _record_to_question(obj: dict[str, Any]) -> TSQuestion:26    return TSQuestion.model_validate(obj)27 28 29def load_json_file(path: Path) -> list[TSQuestion]:30    """Load a single .json file (array or {\"questions\": [...]})."""31    raw = json.loads(path.read_text(encoding="utf-8"))32    records = _parse_records(raw)33    return [_record_to_question(r) for r in records]34 35 36def load_jsonl_file(path: Path) -> list[TSQuestion]:37    """Load newline-delimited JSON; each line must be a full TSQuestion object."""38    out: list[TSQuestion] = []39    for line_no, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1):40        line = line.strip()41        if not line:42            continue43        try:44            obj = json.loads(line)45        except json.JSONDecodeError as e:46            raise ValueError(f"{path}:{line_no}: invalid JSON: {e}") from e47        if not isinstance(obj, dict):48            raise ValueError(f"{path}:{line_no}: expected object per line")49        out.append(_record_to_question(obj))50    return out51 52 53def load_question_banks(54    bank_dir: Path | str | None,55    *,56    domain_order: tuple[str, ...] = DEFAULT_DOMAIN_ORDER,57    explicit_files: list[Path | str] | None = None,58) -> dict[str, list[TSQuestion]]:59    """60    Load per-dataset question pools.61 62    If ``bank_dir`` is set, loads ``<Dataset>_questions.json`` for each domain in63    ``domain_order`` when that file exists, plus any ``*.json`` / ``*.jsonl`` in64    the directory that declare a ``dataset`` field per record (merged lists).65 66    If ``explicit_files`` is set, each file is loaded; records are grouped by67    ``dataset`` field (required for merged files).68    """69    pools: dict[str, list[TSQuestion]] = {d: [] for d in domain_order}70 71    if explicit_files:72        for fp in explicit_files:73            path = Path(fp)74            items = load_jsonl_file(path) if path.suffix.lower() == ".jsonl" else load_json_file(path)75            for q in items:76                if q.dataset not in pools:77                    pools[q.dataset] = []78                pools[q.dataset].append(q)79        return pools80 81    if bank_dir is None:82        return pools83 84    root = Path(bank_dir)85    if not root.is_dir():86        raise NotADirectoryError(f"question_bank_path must be a directory: {root}")87 88    # Per-dataset convention: PSML_questions.json etc.89    for domain in domain_order:90        candidates = [91            root / f"{domain}_questions.json",92            root / f"{domain.lower()}_questions.json",93        ]94        for c in candidates:95            if c.is_file():96                pools[domain].extend(load_json_file(c))97                break98 99    # Any extra json/jsonl with dataset on each row (skip per-dataset files + manifests)100    for path in sorted(root.glob("*.json")) + sorted(root.glob("*.jsonl")):101        if path.name in ("manifest.json", "build_manifest.json"):102            continue103        if any(path.name == f"{d}_questions.json" for d in domain_order):104            continue105        if any(path.name == f"{d.lower()}_questions.json" for d in domain_order):106            continue107        items = load_jsonl_file(path) if path.suffix.lower() == ".jsonl" else load_json_file(path)108        for q in items:109            key = q.dataset110            if key not in pools:111                pools[key] = []112            pools[key].append(q)113 114    return pools115