data.py

"""
Multilingual data prep — Helena (multilingual_model)

Using facebook/belebele as the main dataset — it has exactly our 5 target
languages (Italian, Spanish, Chinese, Russian, Hindi) with 900 MC questions
each. Already in the right format (passage + question + 4 options + answer).

Also loading xcopa for Italian and Chinese as extra data (commonsense MC).

Adapted from trainingPhuc/data.py but much simpler since belebele is
already MC — no need to extract boxed answers from math solutions etc.

Output: sft_train.jsonl, sft_val.jsonl, grpo_train.jsonl
"""

import random
from pathlib import Path

from datasets import load_dataset, Dataset
from transformers import AutoTokenizer


# our 5 target languages + belebele language codes (from the dataset card)
LANGS = {
    "italian": "ita_Latn",
    "spanish": "spa_Latn",
    "chinese": "zho_Hans",
    "russian": "rus_Cyrl",
    "hindi":   "hin_Deva",
}

# system prompt — tells the model to output \boxed{letter}
# keeping it short and language-neutral so it works across all 5 languages
SYSTEM = (
    "Answer the multiple choice question by reasoning step by step. "
    "Write your final answer as \\boxed{A}, \\boxed{B}, \\boxed{C}, etc."
)


def load_belebele(langs=LANGS):
    """Load belebele for all 5 languages.

    Each language has 900 test examples. The dataset has:
    - flores_passage: reading passage in the target language
    - question: question text in the target language
    - mc_answer1-4: the 4 options
    - correct_answer_num: 1-4 (we convert to A-D)
    """
    examples = []
    for lang_name, lang_code in langs.items():
        try:
            ds = load_dataset(
                "facebook/belebele", lang_code,
                split="test", trust_remote_code=True
            )
        except Exception as e:
            print(f"  WARNING: could not load belebele/{lang_code}: {e}")
            continue

        for row in ds:
            answer_letter = "ABCD"[int(row["correct_answer_num"]) - 1]

            # format: passage + question + labeled options
            prompt = (
                f"{row['flores_passage'].strip()}\n\n"
                f"{row['question'].strip()}\n\n"
                f"A) {row['mc_answer1']}\n"
                f"B) {row['mc_answer2']}\n"
                f"C) {row['mc_answer3']}\n"
                f"D) {row['mc_answer4']}"
            )

            examples.append({
                "prompt":   prompt,
                "answer":   answer_letter,
                "language": lang_name,
                "source":   "belebele",
            })

        n = sum(1 for e in examples if e["language"] == lang_name)
        print(f"  belebele/{lang_code} ({lang_name}): {n} examples")

    return examples


def load_xcopa(langs={"italian": "it", "chinese": "zh"}):
    """Extra data from XCOPA for Italian and Chinese.

    XCOPA is commonsense MC with 2 options (cause/effect).
    ~500 examples per language across train/val/test splits.
    Not available for all 5 languages so we only use it where it exists.
    """
    examples = []
    for lang_name, lang_code in langs.items():
        count = 0
        for split in ["train", "validation", "test"]:
            try:
                ds = load_dataset("xcopa", lang_code, split=split)
            except Exception:
                continue

            for row in ds:
                # xcopa has "cause" or "effect" as the question type
                q_suffix = (
                    "What was the cause of this?"
                    if row["question"] == "cause"
                    else "What happened as a result?"
                )
                prompt = (
                    f"{row['premise'].strip()} {q_suffix}\n\n"
                    f"A) {row['choice1']}\n"
                    f"B) {row['choice2']}"
                )
                answer = "AB"[int(row["label"])]

                examples.append({
                    "prompt":   prompt,
                    "answer":   answer,
                    "language": lang_name,
                    "source":   "xcopa",
                })
                count += 1

        print(f"  xcopa/{lang_code} ({lang_name}): {count} examples")

    return examples


def format_for_sft(examples, tokenizer):
    """Format examples as SFT training rows using the chat template.

    Assistant response is just \boxed{letter} — the model will generate
    its own reasoning in the <think> block at inference time (thinking=ON).
    """
    rows = []
    for ex in examples:
        messages = [
            {"role": "system",    "content": SYSTEM},
            {"role": "user",      "content": ex["prompt"]},
            {"role": "assistant", "content": f"\\boxed{{{ex['answer']}}}"},
        ]
        text = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=False
        )
        rows.append({
            "text":     text,
            "prompt":   ex["prompt"],
            "answer":   ex["answer"],
            "language": ex["language"],
            "source":   ex["source"],
        })
    return Dataset.from_list(rows)


def format_for_grpo(examples):
    """Format for GRPOTrainer — prompt only, no assistant turn.

    GRPO samples its own completions and scores them with the reward function.
    Gold answer is kept as a column for the reward function to compare against.
    """
    rows = []
    for ex in examples:
        rows.append({
            "prompt": [
                {"role": "system", "content": SYSTEM},
                {"role": "user",   "content": ex["prompt"]},
            ],
            "answer":   ex["answer"],
            "language": ex["language"],
            "source":   ex["source"],
        })
    return Dataset.from_list(rows)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--model",       default="/shared-ro/models/Qwen/Qwen3-1.7B")
    parser.add_argument("--output-dir",  default="/scratch/multilingual_data")
    parser.add_argument("--no-xcopa",    action="store_true",
                        help="skip loading xcopa (use only belebele)")
    parser.add_argument("--val-frac",    type=float, default=0.05)
    parser.add_argument("--seed",        type=int,   default=42)
    args = parser.parse_args()

    tokenizer = AutoTokenizer.from_pretrained(args.model)
    out = Path(args.output_dir)
    out.mkdir(parents=True, exist_ok=True)

    print("Loading belebele...")
    examples = load_belebele()

    if not args.no_xcopa:
        print("Loading xcopa (Italian + Chinese)...")
        examples += load_xcopa()

    print(f"\nTotal: {len(examples)} examples")

    # shuffle and split
    random.seed(args.seed)
    random.shuffle(examples)

    n_val = int(len(examples) * args.val_frac)
    val_ex   = examples[:n_val]
    train_ex = examples[n_val:]
    print(f"Train: {len(train_ex)}, Val: {len(val_ex)}")

    # save all three splits
    format_for_sft(train_ex, tokenizer).to_json(out / "sft_train.jsonl")
    format_for_sft(val_ex,   tokenizer).to_json(out / "sft_val.jsonl")
    format_for_grpo(train_ex).to_json(out / "grpo_train.jsonl")

    print(f"\nDone. Saved to {out}/")