""" 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 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}/")