"""Add a `reward_model` column to sft_5k.parquet + eval_100.parquet.

verl's reward loop expects each row to have:
    reward_model: {"style": "rule", "ground_truth": <str>}

For OPD with `use_task_rewards=False` the score is unused by the optimizer,
but the loop still calls compute_score() for logging. We store the gold
assistant response as the ground_truth so logging is at least informative.
"""
from __future__ import annotations

from pathlib import Path

import pandas as pd

ROOT = Path("/mnt/local-fast/opd_zt")
DATA = ROOT / "data"


def add_reward_col(pq_path: Path) -> None:
    df = pd.read_parquet(pq_path)
    if "reward_model" in df.columns:
        # Already added.
        sample = df["reward_model"].iloc[0]
        print(f"[skip] {pq_path.name}  already has reward_model, sample={sample!r}")
        return
    df["reward_model"] = [
        {"style": "rule", "ground_truth": str(r)} for r in df["response"]
    ]
    df.to_parquet(pq_path, index=False)
    print(f"[done] {pq_path.name}  rows={len(df)}  sample={df['reward_model'].iloc[0]!r}")


def main() -> None:
    for name in ["sft_5k.parquet", "eval_100.parquet"]:
        add_reward_col(DATA / name)


if __name__ == "__main__":
    main()