Upload 6 files

f908c84 verified 6 months ago

7.19 kB

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

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