data_loader.py

import csv
import json
import os
from pathlib import Path

_HERE = Path(__file__).resolve().parent

DATASETS = {
    "imobench": {
        "csv": _HERE / "data" / "imobench" / "answerbench_v2.csv",
        "task": "math",
    },
    "hle": {
        "csv": _HERE / "data" / "hle" / "text_mc_nomath.csv",
        "task": "mc",
    },
}


def load_queries(dataset, limit=None):
    csv_path = DATASETS[dataset]["csv"]
    if not csv_path.exists():
        hint = "  Run `cd data && python download_hle.py` first." if dataset == "hle" else ""
        raise FileNotFoundError(f"{csv_path} not found.{hint}")
    queries = []
    with open(csv_path) as f:
        for row in csv.DictReader(f):
            queries.append({
                "id": row["Problem ID"],
                "question": row["Problem"].strip(),
                "answer": row["Short Answer"].strip(),
                "extra": {
                    "category": row["Category"],
                    "subcategory": row["Subcategory"],
                    "source": row["Source"],
                },
            })
    return queries[:limit] if limit else queries


def run_dir(dataset, model, n):
    path = _HERE / "runs" / dataset / model / f"predebate_n{n}"
    path.mkdir(parents=True, exist_ok=True)
    return path


def already_done(d):
    return {f[:-5] for f in os.listdir(d) if f.endswith(".json")}


def non_unanimous_qids(predebate_dir, task):
    """Set of query_ids whose round-0 answers are not all identical. Math
    uses parsed_math_answer (fallback boxed_answer); MC uses boxed_answer."""
    out = set()
    for f in sorted(Path(predebate_dir).glob("*.json")):
        d = json.loads(f.read_text())
        ans = []
        for ag in d.get("agents", []):
            sr = ag.get("self_round") or {}
            if task == "math":
                a = sr.get("parsed_math_answer") or sr.get("boxed_answer")
            else:
                a = sr.get("boxed_answer")
            ans.append(a)
        if len({a for a in ans if a is not None}) > 1:
            out.add(d["query_id"])
    return out


def load_predebate(predebate_dir, only_non_unanimous=False):
    """Return [(query, predebate_payload), ...] sorted by query_id.

    `query` mirrors the dict shape used by load_queries (id, question, answer, extra).
    `predebate_payload` is the full per-question JSON written by generate_predebate.py.
    """
    p = Path(predebate_dir)
    files = sorted(p.glob("*.json"))
    out = []
    for f in files:
        with open(f) as fh:
            payload = json.load(fh)
        if only_non_unanimous:
            answers = [ag["self_round"].get("boxed_answer") for ag in payload["agents"]]
            if len(set(answers)) <= 1:
                continue
        query = {
            "id": payload["query_id"],
            "question": payload["question"],
            "answer": payload["ground_truth_answer"],
            "extra": payload.get("extra", {}),
        }
        out.append((query, payload))
    return out