#!/usr/bin/env python3
"""Convert VibeThinker repo math parquets -> flat jsonl the harness eats.

Source rows (verl format):
  prompt:        [{"role":"user","content": "<problem ... \\boxed{}>"}]
  reward_model:  {"ground_truth": "<answer>", "style": "rule"}

Output rows:
  {"id","category","question","answer"}

The prompt already carries the "output the final answer within \\boxed{}"
instruction, so we send it verbatim with NO extra system prompt -- identical
input for every model under test (the controlled-input invariant).
"""
import os, json, glob, argparse
import pandas as pd

DEFAULT_SRC = os.path.expanduser(
    "~/Documents/AI/VibeThinker/eval/math/data")  # adjust on the host

def convert(src, out):
    os.makedirs(out, exist_ok=True)
    total = 0
    for f in sorted(glob.glob(os.path.join(src, "*.parquet"))):
        cat = os.path.splitext(os.path.basename(f))[0]   # aime / aime25 / hmmt25 / gpqa
        df = pd.read_parquet(f)
        rows = []
        for i, r in df.iterrows():
            prompt = r["prompt"]
            # prompt is a numpy array / list of {role,content}
            content = prompt[0]["content"] if not isinstance(prompt, str) else prompt
            gt = r["reward_model"]["ground_truth"]
            rows.append({
                "id": f"{cat}-{i:03d}",
                "category": cat,
                "question": content,
                "answer": str(gt).strip(),
            })
        with open(os.path.join(out, f"{cat}.jsonl"), "w") as fh:
            for row in rows:
                fh.write(json.dumps(row, ensure_ascii=False) + "\n")
        print(f"{cat:<8} {len(rows):>4} items -> {cat}.jsonl")
        total += len(rows)
    print(f"TOTAL {total} items")

if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", default=DEFAULT_SRC)
    ap.add_argument("--out", default=os.path.join(os.path.dirname(__file__), "data"))
    a = ap.parse_args()
    convert(a.src, a.out)