data_utils.py

from datasets import load_dataset, concatenate_datasets, Dataset
from functools import partial
from transformers import AutoTokenizer
import random
import json
import re
import hashlib
from dataclasses import dataclass
from typing import Callable

def example_map_fn(example, idx, process_fn, data_source, ability, split):
    question, solution = process_fn(example)
    data = {
        "data_source": data_source,
        "prompt": [{"role": "user", "content": question}],
        "question": question,
        "ability": ability,
        "reward_model": {"style": "rule", "ground_truth": solution},
        "extra_info": {"question": question, "split": split, "index": idx, "need_tools_kwargs": True,
                    "tools_kwargs": {"ground_truth": solution, "split": split},
                    },
        "agent_name": "tool_agent",
    }
    return data

def build_verl_parquet_openr1_bigmath_oneshot(subset="level_5", max_unique_prompts=1024, max_train_size=1024, seed=42):
    data_source = "open-r1/Big-Math-RL-Verified-Processed"
    ability = "math"

    train_ds = load_dataset(data_source, subset, split="train").shuffle(seed=seed, load_from_cache_file=False)

    if max_unique_prompts and len(train_ds) > max_unique_prompts:
        train_ds = train_ds.select(range(max_unique_prompts))

    if max_train_size is not None:
        n = len(train_ds)
        if n < max_train_size:
            repeats, remainder = divmod(max_train_size, n)
            idx = list(range(n)) * repeats
            idx += list(range(remainder))
            train_ds = train_ds.select(idx)
        else:
            train_ds = train_ds.select(range(max_train_size))

    def process_openr1_bigmath(example):
        user_prompt = example.get("prompt", "")
        return user_prompt, example.get("solution")

    train_map_fn = partial(
        example_map_fn, process_fn=process_openr1_bigmath, data_source=data_source, ability=ability, split="train"
    )
    train_ds = train_ds.map(train_map_fn, with_indices=True, remove_columns=train_ds.column_names, load_from_cache_file=False)
    return train_ds

def build_aime2024_dataset():
    def process_aime2024(example):
        problem = example["Problem"]
        # Force parseable final line
        problem = (
            problem
            + "\n\nGive ONLY the final answer on the last line in exactly one of these formats:\n"
            r"Answer: \boxed{<integer>}"
            "\n"
        )
        return problem, str(example["Answer"])

    data_source = "Maxwell-Jia/AIME_2024"
    print(f"Loading the {data_source} dataset from huggingface...", flush=True)
    dataset = load_dataset(data_source, split="train")
    map_fn = partial(
        example_map_fn, process_fn=process_aime2024, data_source=data_source, ability="Math", split="test"
    )
    dataset = dataset.map(map_fn, with_indices=True, remove_columns=dataset.column_names, load_from_cache_file=False)
    return dataset


def build_gpqa_diamond_dataset():
    import random

    GPQA_QUERY_TEMPLATE = (
        "Answer the following multiple choice question. The last line of your response should be of the following "
        "format: 'Answer: $LETTER' (without quotes) where LETTER is one of ABCD. Think step by step before "
        "answering.\n\n{Question}\n\nA) {A}\nB) {B}\nC) {C}\nD) {D}"
    )

    def process_gpqa_diamond(example):
        choices = [example["Incorrect Answer 1"], example["Incorrect Answer 2"], example["Incorrect Answer 3"]]
        random.shuffle(choices)
        gold_index = random.randint(0, 3)
        choices.insert(gold_index, example["Correct Answer"])
        query_prompt = GPQA_QUERY_TEMPLATE.format(
            A=choices[0], B=choices[1], C=choices[2], D=choices[3], Question=example["Question"]
        )
        gold_choice = "ABCD"[gold_index]
        return query_prompt, gold_choice

    data_source = "Idavidrein/gpqa"
    print(f"Loading the {data_source} dataset from huggingface...", flush=True)

    dataset = load_dataset(data_source, "gpqa_diamond", split="train")
    map_fn = partial(
        example_map_fn, process_fn=process_gpqa_diamond, data_source=data_source, ability="Science", split="test"
    )
    dataset = dataset.map(map_fn, with_indices=True, remove_columns=dataset.column_names, load_from_cache_file=False,)
    return dataset


def build_cnmo2024_dataset():
    def process_cnmo2024(example):
        return example["question"], example["answer"]

    data_source = "opencompass/LiveMathBench"
    print(f"Loading the {data_source} dataset from huggingface...", flush=True)

    dataset_en = load_dataset(data_source, "v202412_CNMO_en", split="test")
    map_fn_en = partial(
        example_map_fn, process_fn=process_cnmo2024, data_source="opencompass/cnmo2024_en", ability="Math", split="test"
    )
    dataset_en = dataset_en.map(map_fn_en, with_indices=True, remove_columns=dataset_en.column_names, load_from_cache_file=False,)

    dataset_zh = load_dataset(data_source, "v202412_CNMO_cn", split="test")
    map_fn_zh = partial(
        example_map_fn, process_fn=process_cnmo2024, data_source="opencompass/cnmo2024_zh", ability="Math", split="test"
    )
    dataset_zh = dataset_zh.map(map_fn_zh, with_indices=True, remove_columns=dataset_zh.column_names, load_from_cache_file=False,)

    dataset = concatenate_datasets([dataset_en, dataset_zh])
    return dataset


def build_livecodebench_dataset():
    import base64
    import json
    import pickle
    import zlib

    def process_livecodebench(example):
        # Construct Query Prompt
        # From https://github.com/LiveCodeBench/LiveCodeBench/blob/998c52d394b836f15fff3b9a29866191108ff81b/lcb_runner/prompts/code_generation.py#L140
        query_prompt = (
            f"You will be given a question (problem specification) and will generate a correct Python program "
            f"that matches the specification and passes all tests.\n\nQuestion: {example['question_content']}\n\n"
        )
        if example["starter_code"]:
            query_prompt += (
                f"You will use the following starter code to write the solution to the problem and enclose your "
                f"code within delimiters.\n```python\n{example['starter_code']}\n```"
            )
        else:
            query_prompt += (
                "Read the inputs from stdin solve the problem and write the answer to stdout (do not directly test "
                "on the sample inputs). Enclose your code within delimiters as follows. Ensure that when the python "
                "program runs, it reads the inputs, runs the algorithm and writes output to STDOUT."
                "```python\n# YOUR CODE HERE\n```"
            )

        # Construct test cases
        public_test_cases = json.loads(example["public_test_cases"])
        try:
            private_test_cases = json.loads(example["private_test_cases"])
        except Exception as e:
            print(f"Error loading private test cases: {e}")
            private_test_cases = json.loads(
                pickle.loads(zlib.decompress(base64.b64decode(example["private_test_cases"].encode("utf-8"))))
            )
        full_test_cases = public_test_cases + private_test_cases

        metadata = json.loads(example["metadata"])
        test_cases = {
            "inputs": [t["input"] for t in full_test_cases],
            "outputs": [t["output"] for t in full_test_cases],
            "fn_name": metadata.get("func_name", None),
        }
        text_cases_compressed = base64.b64encode(zlib.compress(pickle.dumps(json.dumps(test_cases)))).decode("utf-8")
        return query_prompt, text_cases_compressed

    data_source = "livecodebench/code_generation_lite"
    print(f"Loading the {data_source} dataset from huggingface...", flush=True)
    dataset = load_dataset(data_source, split="test")
    # R1 Evaluation use LiveCodeBench 24.08-25.01
    dataset = dataset.filter(lambda line: "2024-08-00T00:00:00" <= line["contest_date"] < "2025-01-00T00:00:00")
    map_fn = partial(
        example_map_fn, process_fn=process_livecodebench, data_source=data_source, ability="Code", split="test"
    )

    dataset = dataset.map(map_fn, with_indices=True, remove_columns=dataset.column_names, num_proc=8, load_from_cache_file=False)
    return dataset


### Multi problem prompt dataset building functions

def pack_problems_greedy(
    all_token_counts: list[int],
    header_tokens: int,
    max_prompt_tokens: int,
    problems_per_prompt: int,
    num_multi_prompts: int,
    seed: int = 42,
) -> list[list[int]]:
    """
    Greedily pack problems into prompts respecting BOTH:
      - token budget:  header + sum(problem_tokens) <= max_prompt_tokens
      - count budget:  len(batch) <= problems_per_prompt

    Returns list of batches, each batch is a list of indices.
    """
    rng = random.Random(seed)
    n = len(all_token_counts)
    indices = list(range(n))
    batches = []

    for _ in range(num_multi_prompts):
        rng.shuffle(indices)
        batch = []
        seen = set()
        running_tokens = header_tokens

        for idx in indices:
            if len(batch) >= problems_per_prompt:
                break
            if idx in seen:
                continue

            numbering_overhead = 5
            candidate_cost = running_tokens + all_token_counts[idx] + numbering_overhead

            if candidate_cost > max_prompt_tokens:
                continue

            batch.append(idx)
            seen.add(idx)
            running_tokens = candidate_cost

        if not batch:
            shortest = min(range(n), key=lambda i: all_token_counts[i])
            batch = [shortest]

        batches.append(batch)

    return batches

def _make_prompt_id(batch_indices: list[int], data_source: str) -> str:
    """Stable hash of batch composition for exclusion-rule enforcement."""
    key = f"{data_source}:{sorted(batch_indices)}"
    return hashlib.sha256(key.encode()).hexdigest()[:16]

def _setup_tokenizer(tokenizer_name_or_path: str, max_prompt_length: int):
    """Load tokenizer, compute chat overhead, return (tokenizer, max_prompt_tokens)."""
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)

    chat_overhead = 0
    if hasattr(tokenizer, "apply_chat_template"):
        empty_chat = tokenizer.apply_chat_template(
            [{"role": "user", "content": ""}],
            tokenize=True,
            add_generation_prompt=True,
        )
        bare_empty = tokenizer.encode("", add_special_tokens=False)
        chat_overhead = len(empty_chat) - len(bare_empty)

    max_prompt_tokens = max_prompt_length - chat_overhead
    print(f"Token budget: {max_prompt_length} - {chat_overhead} chat overhead = {max_prompt_tokens} prompt tokens")

    assert max_prompt_tokens > 100, (
        f"max_prompt_tokens={max_prompt_tokens} is too small. Increase max_prompt_length."
    )
    return tokenizer, max_prompt_tokens

def _tokenize_block(tokenizer, question: str, idx: int = 1) -> int:
    """Token count for a single '## Problem {idx}\\n{question}\\n\\n' block."""
    block = f"## Problem {idx}\n{question.strip()}\n\n"
    return len(tokenizer.encode(block, add_special_tokens=False))

def _estimate_header_tokens(tokenizer, build_prompt_fn, dummy_count: int = 10, **prompt_kwargs) -> int:
    """Tokenize header by subtracting dummy problem blocks from full prompt."""
    header = build_prompt_fn(["dummy"] * dummy_count, dummy_count, **prompt_kwargs)
    problems_only = "".join(f"## Problem {i}\ndummy\n\n" for i in range(1, dummy_count + 1))
    return max(
        len(tokenizer.encode(header, add_special_tokens=False))
        - len(tokenizer.encode(problems_only, add_special_tokens=False)),
        0,
    )

def _pack_and_build(
    all_questions: list[str],
    tokenizer,
    max_prompt_tokens: int,
    problems_per_prompt: int,
    max_train_size: int,
    seed: int,
    build_prompt_fn: Callable,
    build_prompt_kwargs: dict,
    build_record_fn: Callable,
    build_record_kwargs: dict,
    data_source: str,
) -> Dataset:
    """
    Shared core: pre-tokenize → filter oversized → pack → build records → verify → stats.

    Args:
        all_questions:       formatted question strings
        tokenizer:           loaded tokenizer
        max_prompt_tokens:   token cap for prompt content (after chat overhead)
        problems_per_prompt: max problems per prompt
        max_train_size:      number of multi-problem prompts to emit
        seed:                reproducibility
        build_prompt_fn:     (questions, num_problems, **kwargs) -> str
        build_prompt_kwargs: extra kwargs for build_prompt_fn (e.g. language)
        build_record_fn:     (batch_indices, all_questions, prompt_text, **kwargs) -> dict
        build_record_kwargs: extra kwargs for build_record_fn
        data_source:         for prompt_id generation
    """
    num_raw = len(all_questions)

    # ── pre-tokenize ──
    all_token_counts = [_tokenize_block(tokenizer, q, 1) for q in all_questions]
    header_tokens = _estimate_header_tokens(
        tokenizer, build_prompt_fn, dummy_count=problems_per_prompt, **build_prompt_kwargs
    )

    avg_tok = sum(all_token_counts) / len(all_token_counts)
    max_tok = max(all_token_counts)
    min_tok = min(all_token_counts)
    print(f"Problem token stats: min={min_tok}, avg={avg_tok:.0f}, max={max_tok}")
    print(f"Header tokens: ~{header_tokens}")
    print(f"Theoretical max problems/prompt at avg: "
          f"{(max_prompt_tokens - header_tokens) // avg_tok:.0f}")

    # ── drop oversized ──
    solo_budget = max_prompt_tokens - header_tokens
    oversized = {i for i, t in enumerate(all_token_counts) if t > solo_budget}
    if oversized:
        print(f"⚠ Dropping {len(oversized)} problems exceeding solo budget ({solo_budget} tok)")
        keep = [i for i in range(num_raw) if i not in oversized]
        # Return the keep mask so callers can filter their parallel arrays
        all_questions_filtered = [all_questions[i] for i in keep]
        all_token_counts = [all_token_counts[i] for i in keep]
    else:
        keep = list(range(num_raw))
        all_questions_filtered = all_questions

    # ── pack ──
    num_multi_prompts = max_train_size or (len(all_questions_filtered) // problems_per_prompt)

    batches = pack_problems_greedy(
        all_token_counts=all_token_counts,
        header_tokens=header_tokens,
        max_prompt_tokens=max_prompt_tokens,
        problems_per_prompt=problems_per_prompt,
        num_multi_prompts=num_multi_prompts,
        seed=seed,
    )

    # ── build records + verify ──
    records = []
    clipped = 0
    for batch_idx_local in batches:
        # Map local indices (post-filter) back to original indices
        batch_idx_orig = [keep[i] for i in batch_idx_local]

        questions_batch = [all_questions_filtered[i] for i in batch_idx_local]
        prompt_text = build_prompt_fn(questions_batch, len(questions_batch), **build_prompt_kwargs)

        actual_tokens = len(tokenizer.encode(prompt_text, add_special_tokens=False))

        # Trim from the end if over budget
        while len(batch_idx_local) > 1 and actual_tokens > max_prompt_tokens:
            clipped += 1 if len(batch_idx_local) == len(batch_idx_orig) else 0
            batch_idx_local = batch_idx_local[:-1]
            batch_idx_orig = batch_idx_orig[:-1]
            questions_batch = [all_questions_filtered[i] for i in batch_idx_local]
            prompt_text = build_prompt_fn(questions_batch, len(questions_batch), **build_prompt_kwargs)
            actual_tokens = len(tokenizer.encode(prompt_text, add_special_tokens=False))

        prompt_id = _make_prompt_id(batch_idx_orig, data_source)

        record = build_record_fn(
            batch_indices=batch_idx_orig,
            prompt_text=prompt_text,
            prompt_id=prompt_id,
            **build_record_kwargs,
        )
        records.append(record)

    if clipped:
        print(f"⚠ {clipped} prompts needed post-hoc trimming")

    # ── final stats ──
    pcounts = [r["extra_info"]["num_problems"] for r in records]
    print(f"\nFinal dataset: {len(records)} multi-problem prompts")
    print(f"Problems per prompt: min={min(pcounts)}, avg={sum(pcounts)/len(pcounts):.1f}, max={max(pcounts)}")

    return Dataset.from_list(records)

def build_math_multi_prompt(questions: list[str], num_problems: int) -> str:
    header = (
        f"You are given {num_problems} math problems below. "
        "Solve as many as you can, one by one, in order. "
        "For each problem, wrap your final answer in \\boxed{{}}.\n"
        "Use the format:\n"
        "## Problem K\n"
        "<your work>\n"
        "**Answer K:** \\boxed{<answer>}\n\n"
        "If you cannot solve a problem, write **Answer:** \\boxed{SKIP} and move on.\n"
        "Solve as many problems as possible before you run out of space.\n\n"
        "---\n\n"
    )
    body = "".join(f"## Problem {i}\n{q.strip()}\n\n" for i, q in enumerate(questions, 1))
    return header + body

def _math_build_record(
    batch_indices: list[int],
    prompt_text: str,
    prompt_id: str,
    all_questions: list[str],
    all_solutions: list[str],
    data_source: str,
    ability: str,
    split: str,
):
    questions = [all_questions[i] for i in batch_indices]
    solutions = [all_solutions[i] for i in batch_indices]
    return {
        "data_source": data_source,
        "prompt": [{"role": "user", "content": prompt_text}],
        "question": prompt_text,
        "ability": ability,
        "reward_model": {
            "style": "rule",
            "ground_truth": json.dumps(solutions),
        },
        "extra_info": {
            "questions": questions,
            "solutions": solutions,
            "num_problems": len(questions),
            "prompt_id": prompt_id,
            "split": split,
            "index": batch_indices,
            "need_tools_kwargs": True,
            "tools_kwargs": {
                "ground_truths": solutions,
                "split": split,
            },
        },
        "agent_name": "tool_agent",
    }

def build_verl_parquet_openr1_bigmath_multi(
    subset="level_5",
    max_unique_prompts=1024,
    max_train_size=1024,
    problems_per_prompt=5,
    max_prompt_length=8192,
    tokenizer_name_or_path="Qwen/Qwen2.5-7B",
    seed=42,
) -> Dataset:
    data_source = "open-r1/Big-Math-RL-Verified-Processed"
    ability = "math"

    tokenizer, max_prompt_tokens = _setup_tokenizer(tokenizer_name_or_path, max_prompt_length)

    raw_ds = load_dataset(data_source, subset, split="train")
    raw_ds = raw_ds.shuffle(seed=seed, load_from_cache_file=False)
    if max_unique_prompts and len(raw_ds) > max_unique_prompts:
        raw_ds = raw_ds.select(range(max_unique_prompts))

    all_questions = [ex.get("prompt", "") for ex in raw_ds]
    all_solutions = [ex.get("solution", "") for ex in raw_ds]

    return _pack_and_build(
        all_questions=all_questions,
        tokenizer=tokenizer,
        max_prompt_tokens=max_prompt_tokens,
        problems_per_prompt=problems_per_prompt,
        max_train_size=max_train_size,
        seed=seed,
        build_prompt_fn=build_math_multi_prompt,
        build_prompt_kwargs={},
        build_record_fn=_math_build_record,
        build_record_kwargs=dict(
            all_questions=all_questions,
            all_solutions=all_solutions,
            data_source=data_source,
            ability=ability,
            split="train",
        ),
        data_source=data_source,
    )

def _format_cf_problem(example: dict) -> str:
    """Turn one row of open-r1/codeforces into a self-contained problem string."""
    parts = []

    title = example.get("title", "")
    if title:
        parts.append(f"**{title}**")

    tl = example.get("time_limit")
    ml = example.get("memory_limit")
    if tl or ml:
        constraints = []
        if tl:
            constraints.append(f"time limit: {tl}s")
        if ml:
            constraints.append(f"memory limit: {int(ml)}MB")
        parts.append(f"({', '.join(constraints)})")

    desc = example.get("description", "")
    if desc:
        parts.append(desc.strip())

    inp_fmt = example.get("input_format", "")
    if inp_fmt:
        parts.append(f"**Input**\n{inp_fmt.strip()}")

    out_fmt = example.get("output_format", "")
    if out_fmt:
        parts.append(f"**Output**\n{out_fmt.strip()}")

    interaction_fmt = example.get("interaction_format", "")
    if interaction_fmt:
        parts.append(f"**Interaction**\n{interaction_fmt.strip()}")

    examples = example.get("examples", [])
    if examples:
        ex_lines = ["**Examples**"]
        for i, ex in enumerate(examples, 1):
            ex_lines.append(f"Input {i}:\n```\n{ex['input'].strip()}\n```")
            ex_lines.append(f"Output {i}:\n```\n{ex['output'].strip()}\n```")
        parts.append("\n".join(ex_lines))

    note = example.get("note", "")
    if note:
        parts.append(f"**Note**\n{note.strip()}")

    return "\n\n".join(parts)

def _extract_ground_truth(example: dict) -> dict:
    """Pack everything the verifier needs into a single dict."""
    tests = [{"input": t["input"], "output": t["output"]} for t in (example.get("official_tests") or [])]
    gen_tests = example.get("generated_tests") or []
    return {
        "id": example.get("id", ""),
        "tests": tests,
        "generated_checker": example.get("generated_checker"),
        "time_limit": example.get("time_limit"),
        "memory_limit": example.get("memory_limit"),
        "input_mode": example.get("input_mode", "stdio"),
        "generated_tests_available": len(gen_tests) if isinstance(gen_tests, list) else int(gen_tests or 0),
    }

def build_cf_multi_prompt(questions: list[str], num_problems: int, language: str = "python") -> str:
    lang_label = "Python 3" if language == "python" else "C++17"
    header = (
        f"You are given {num_problems} competitive programming problems below.\n"
        f"Solve as many as you can, one by one, in order.\n"
        f"For each problem, provide a complete, correct **{lang_label}** solution.\n\n"
        "Use the format:\n"
        "## Problem K\n"
        "<your reasoning>\n"
        f"```{language}\n<your code>\n```\n\n"
        "If you cannot solve a problem, write `SKIP` and move on.\n"
        "Your solutions must read from stdin and write to stdout unless stated otherwise.\n"
        "Solve as many problems as possible before you run out of space.\n\n"
        "---\n\n"
    )
    body = "".join(f"## Problem {i}\n{q.strip()}\n\n" for i, q in enumerate(questions, 1))
    return header + body

def _cf_build_record(
    batch_indices: list[int],
    prompt_text: str,
    prompt_id: str,
    all_questions: list[str],
    all_ground_truths: list[dict],
    all_ids: list[str],
    data_source: str,
    ability: str,
    split: str,
    language: str,
):
    questions = [all_questions[i] for i in batch_indices]
    gts = [all_ground_truths[i] for i in batch_indices]
    ids = [all_ids[i] for i in batch_indices]
    return {
        "data_source": data_source,
        "prompt": [{"role": "user", "content": prompt_text}],
        "question": prompt_text,
        "ability": ability,
        "reward_model": {
            "style": "rule",
            "ground_truth": json.dumps(gts),
        },
        "extra_info": {
            "questions": questions,
            "problem_ids": ids,
            "num_problems": len(questions),
            "language": language,
            "prompt_id": prompt_id,
            "split": split,
            "index": batch_indices,
            "need_tools_kwargs": True,
            "tools_kwargs": {
                "ground_truths": gts,
                "language": language,
                "split": split,
            },
        },
        "agent_name": "tool_agent",
    }

def build_verl_parquet_codeforces_multi(
    subset: str = "verifiable",
    rating_range: tuple[int, int] | None = (1400, 1600),
    tags: list[str] | None = None,
    problems_per_prompt: int = 15,
    max_unique_problems: int | None = None,
    max_train_size: int = 256,
    max_prompt_length: int = 16384,
    language: str = "python",
    tokenizer_name_or_path: str = "Qwen/Qwen2.5-7B",
    exclude_interactive: bool = True,
    seed: int = 42,
) -> Dataset:
    """
    Build a verl-compatible dataset where each row is a multi-problem
    codeforces prompt (batch of `problems_per_prompt` problems).
    """
    data_source = "open-r1/codeforces"
    ability = "code"

    tokenizer, max_prompt_tokens = _setup_tokenizer(tokenizer_name_or_path, max_prompt_length)

    # ── load & filter ──
    print(f"Loading {data_source} ({subset}) …")
    raw_ds = load_dataset(data_source, subset, split="train")
    print(f"  loaded {len(raw_ds)} problems")

    if exclude_interactive:
        raw_ds = raw_ds.filter(lambda ex: not ex.get("interaction_format"), load_from_cache_file=False)
        print(f"  after dropping interactive: {len(raw_ds)}")

    if rating_range is not None:
        lo, hi = rating_range
        raw_ds = raw_ds.filter(
            lambda ex: ex.get("rating") is not None and lo <= int(ex["rating"]) <= hi,
            load_from_cache_file=False,
        )
        print(f"  after rating filter [{lo}, {hi}]: {len(raw_ds)}")

    if tags is not None:
        tag_set = set(t.lower() for t in tags)
        raw_ds = raw_ds.filter(
            lambda ex: tag_set.issubset(set(t.lower() for t in (ex.get("tags") or []))),
            load_from_cache_file=False,
        )
        print(f"  after tag filter {tags}: {len(raw_ds)}")

    raw_ds = raw_ds.filter(
        lambda ex: len(ex.get("official_tests") or []) > 0,
        load_from_cache_file=False,
    )
    print(f"  after requiring ≥1 test: {len(raw_ds)}")

    raw_ds = raw_ds.shuffle(seed=seed, load_from_cache_file=False)
    if max_unique_problems and len(raw_ds) > max_unique_problems:
        raw_ds = raw_ds.select(range(max_unique_problems))

    # ── format ──
    all_questions = [_format_cf_problem(ex) for ex in raw_ds]
    all_ground_truths = [_extract_ground_truth(ex) for ex in raw_ds]
    all_ids = [ex.get("id", str(i)) for i, ex in enumerate(raw_ds)]
    print(f"  problem pool: {len(all_questions)}")

    return _pack_and_build(
        all_questions=all_questions,
        tokenizer=tokenizer,
        max_prompt_tokens=max_prompt_tokens,
        problems_per_prompt=problems_per_prompt,
        max_train_size=max_train_size,
        seed=seed,
        build_prompt_fn=build_cf_multi_prompt,
        build_prompt_kwargs=dict(language=language),
        build_record_fn=_cf_build_record,
        build_record_kwargs=dict(
            all_questions=all_questions,
            all_ground_truths=all_ground_truths,
            all_ids=all_ids,
            data_source=data_source,
            ability=ability,
            split="train",
            language=language,
        ),
        data_source=data_source,
    )