code/prepare_gsm_sft.py

"""
Convert the Arabic GSM8K dataset into the SFT schema train_reasoning.py already reads
({instruction, reasoning, answer}), so v2 trains in the same ChatML + <think> format as v1.

The final answer is kept as the bare numeral: it matches the source's intent and makes
exact-match evaluation trivial.
"""
import json
import random
from pathlib import Path

import pyarrow.parquet as pq

SRC = Path("out_gsm/gsm8k_reasoning_ar.parquet")
OUT = Path("data_gsm_sft")
EVAL_N = 2000
SEED = 42


def main():
    d = pq.read_table(SRC).to_pydict()
    rows = [
        {"instruction": q, "reasoning": t, "answer": a}
        for q, t, a in zip(d["question"], d["thinking"], d["answer"])
    ]
    random.Random(SEED).shuffle(rows)
    eval_rows, train_rows = rows[:EVAL_N], rows[EVAL_N:]

    OUT.mkdir(exist_ok=True)
    for name, split in (("train", train_rows), ("eval", eval_rows)):
        with open(OUT / f"{name}.jsonl", "w", encoding="utf-8") as fh:
            for r in split:
                fh.write(json.dumps(r, ensure_ascii=False) + "\n")
        print(f"[+] {name}: {len(split)} -> {OUT / f'{name}.jsonl'}")


if __name__ == "__main__":
    main()