dataloaders/math.py

from pathlib import Path
import re
import numpy as np
import torch
from datasets import Features, Value, load_dataset, load_from_disk, DatasetDict, concatenate_datasets, disable_caching, enable_caching
from torch.utils.data import DataLoader
from torch_utils.distributed import get_rank, get_world_size
from dataloaders.sampler import InfiniteSampler
from evaluate.grader import math_equal
from evaluate.parser import extract_answer, strip_string


ACDIR_REFERENCE_FEATURES = Features(
    {
        "prompt_id": Value("string"),
        "source_dataset": Value("string"),
        "subject": Value("string"),
        "difficulty": Value("string"),
        "level": Value("string"),
        "problem": Value("string"),
        "solution": Value("string"),
        "canonical_final_answer": Value("string"),
        "answer_parser_type": Value("string"),
    }
)


PROFILE_COLUMNS = {
    "dataset_index",
    "baseline_correct",
    "parse_failed",
    "baseline_response",
    "extracted_answer",
    "extracted_response",
    "problem_chars",
    "prompt_chars",
    "response_chars",
    "profile_bucket",
    "profile_job_id",
    "profile_backend",
    "profile_fast_mode",
    "profile_steps",
    "profile_gen_length",
    "profile_block_length",
}


def load_dataset_split(local_path: str, split: str, data_dir: str = None):
    path = Path(local_path)
    if path.is_dir():
        is_saved = (path / "dataset_dict.json").exists() or (path / "state.json").exists()
        if is_saved:
            ds = load_from_disk(str(path))
            if isinstance(ds, DatasetDict):
                return ds[split]
            return ds
    if data_dir:
        return load_dataset(local_path, split=split, data_dir=data_dir)
    return load_dataset(local_path, split=split)



def collate_fn_math(batch,):
    problems = []
    answers = []
    levels = []
    instruct = r"(Please put the final answer in \boxed{} tag, i.e. $\boxed{answer here}$)"
    for item in batch:
        problems.append(item['problem'] + instruct)
        answers.append(item['solution'])
        levels.append(item['level'])
    
    return {
        "problems": problems, 
        "answers": answers,
        "levels": levels,
    }
    

def collate_fn_gsm8k(batch,):
    problems = []
    answers = []
    for item in batch:
        problems.append(item['question'])
        answers.append(item['answer'])

    return {
        "problems": problems, 
        "answers": answers
    }


def _first_present(item, keys, default=""):
    for key in keys:
        if key in item and item[key] is not None:
            value = item[key]
            if isinstance(value, (list, tuple)):
                if value:
                    return str(value[0])
                continue
            return str(value)
    return default


def normalize_math_record(item, source_dataset="unknown"):
    problem = _first_present(
        item,
        (
            "problem",
            "question",
            "prompt",
            "input",
            "instruction",
            "query",
            "original_question",
            "cleaned_problem",
        ),
    )
    if "messages" in item and not problem:
        try:
            messages = item["messages"]
            if messages and isinstance(messages, list):
                problem = str(messages[0].get("content", ""))
        except Exception:
            problem = ""
    solution = _first_present(
        item,
        (
            "solution",
            "answer",
            "output",
            "response",
            "text",
            "thought",
            "generation",
        ),
    )
    canonical_answer = _first_present(item, ("canonical_final_answer", "final_answer"), default="")
    level = _first_present(item, ("level", "difficulty", "problem_type", "source"), default="")
    subject = _first_present(item, ("subject", "type", "problem_type"), default="")
    prompt_id = _first_present(item, ("prompt_id", "uuid", "id"), default="")
    if not prompt_id:
        prompt_id = str(abs(hash((source_dataset, problem[:256], (canonical_answer or solution)[:128]))))
    return {
        "prompt_id": str(prompt_id),
        "source_dataset": str(item.get("source_dataset", source_dataset)),
        "subject": str(subject),
        "difficulty": str(level),
        "level": str(level),
        "problem": str(problem),
        "solution": str(solution if solution else canonical_answer),
        "canonical_final_answer": str(canonical_answer),
        "answer_parser_type": "math",
    }


def collate_fn_acdir_reference(batch):
    problems = []
    answers = []
    solutions = []
    answer_is_canonical = []
    sources = []
    levels = []
    prompt_ids = []
    for item in batch:
        norm = normalize_math_record(item, item.get("source_dataset", "unknown"))
        instruct = r"(Please put the final answer in \boxed{} tag, i.e. $\boxed{answer here}$)"
        problem = norm["problem"]
        if "\\boxed" not in problem and "boxed" not in problem.lower():
            problem = problem + instruct
        problems.append(problem)
        answer = norm["canonical_final_answer"] or norm["solution"]
        answers.append(answer)
        solutions.append(norm["solution"])
        answer_is_canonical.append(bool(norm["canonical_final_answer"]))
        sources.append(norm["source_dataset"])
        levels.append(norm["level"])
        prompt_ids.append(norm["prompt_id"])
    return {
        "problems": problems,
        "answers": answers,
        "solutions": solutions,
        "answer_is_canonical": answer_is_canonical,
        "sources": sources,
        "levels": levels,
        "prompt_ids": prompt_ids,
    }


def collate_fn_acdir_profiled_reference(batch):
    out = collate_fn_acdir_reference(batch)
    out.update(
        {
            "dataset_indices": torch.tensor(
                [int(item.get("dataset_index", -1)) for item in batch],
                dtype=torch.long,
            ),
            "profile_baseline_correct": torch.tensor(
                [bool(item.get("baseline_correct", False)) for item in batch],
                dtype=torch.bool,
            ),
            "profile_parse_failed": torch.tensor(
                [bool(item.get("parse_failed", False)) for item in batch],
                dtype=torch.bool,
            ),
            "profile_buckets": [str(item.get("profile_bucket", "")) for item in batch],
            "profile_job_ids": [str(item.get("profile_job_id", "")) for item in batch],
            "profile_baseline_responses": [str(item.get("baseline_response", "")) for item in batch],
            "profile_extracted_answers": [str(item.get("extracted_answer", "")) for item in batch],
            "profile_extracted_responses": [str(item.get("extracted_response", "")) for item in batch],
            "profile_response_chars": torch.tensor(
                [int(item.get("response_chars", 0) or 0) for item in batch],
                dtype=torch.long,
            ),
        }
    )
    return out


class ProfileBucketSampler(torch.utils.data.Sampler):
    def __init__(
        self,
        dataset,
        bucket_weights,
        *,
        rank=0,
        num_replicas=1,
        seed=0,
        exclude_parse_failed=True,
    ):
        assert len(dataset) > 0
        super().__init__()
        self.dataset = dataset
        self.rank = int(rank)
        self.num_replicas = int(num_replicas)
        self.seed = int(seed)
        self.exclude_parse_failed = bool(exclude_parse_failed)

        weights = dict(bucket_weights or {})
        if not weights:
            weights = {
                "hard_wrong_mathlike": 0.45,
                "medium_wrong_mathlike": 0.20,
                "right_safety": 0.30,
                "random_audit": 0.05,
            }
        self.bucket_indices = self._build_bucket_indices()
        names = []
        probs = []
        for name, weight in weights.items():
            weight = float(weight)
            if weight <= 0:
                continue
            if name not in self.bucket_indices or len(self.bucket_indices[name]) <= 0:
                continue
            names.append(name)
            probs.append(weight)
        if not names:
            raise ValueError(f"No non-empty profile buckets for weights: {weights}")
        probs = np.asarray(probs, dtype=np.float64)
        probs = probs / probs.sum()
        self.bucket_names = names
        self.bucket_probs = probs

    def _build_bucket_indices(self):
        buckets = {}
        actual = list(self.dataset["profile_bucket"])
        correct = list(self.dataset["baseline_correct"])
        parse_failed = list(self.dataset["parse_failed"])
        all_nonparse = []
        wrong = []
        right = []
        for idx, bucket in enumerate(actual):
            is_parse_failed = bool(parse_failed[idx])
            if self.exclude_parse_failed and is_parse_failed:
                continue
            bucket = str(bucket)
            buckets.setdefault(bucket, []).append(idx)
            all_nonparse.append(idx)
            if bool(correct[idx]):
                right.append(idx)
            else:
                wrong.append(idx)
        buckets["random_audit"] = all_nonparse
        buckets["baseline_wrong"] = wrong
        buckets["baseline_right"] = right
        return {key: np.asarray(value, dtype=np.int64) for key, value in buckets.items()}

    def __iter__(self):
        rng = np.random.RandomState(self.seed + self.rank * 1000003)
        while True:
            bucket_name = self.bucket_names[int(rng.choice(len(self.bucket_names), p=self.bucket_probs))]
            indices = self.bucket_indices[bucket_name]
            yield int(indices[int(rng.randint(0, len(indices)))])

    def __len__(self):
        return len(self.dataset)


def try_get_level(level: str, default: int = 5):
    try:
        return int(level.split()[-1])
    except:
        return default


def normalize_reference_answer(answer: str, *, canonical: bool = False) -> str:
    answer = "" if answer is None else str(answer).strip()
    if not answer:
        return ""
    if not canonical:
        return extract_answer(answer)
    lower = answer.lower()
    if "####" in answer or "boxed" in lower or "final answer is" in lower or "he answer is" in lower or "答案是" in answer:
        return extract_answer(answer, skip_unit=True)
    return strip_string(answer, skip_unit=True)


def has_parseable_reward_reference(item, max_chars: int = 256) -> bool:
    ref = item.get("canonical_final_answer") or item.get("solution") or ""
    extracted = normalize_reference_answer(ref, canonical=bool(item.get("canonical_final_answer")))
    return bool(str(extracted).strip()) and len(str(extracted)) <= int(max_chars)


def is_proof_like_problem(item) -> bool:
    problem = str(item.get("problem", "") or "").strip().lower()
    problem = re.sub(r"\s+", " ", problem)
    if not problem:
        return False
    return (
        bool(re.search(r"\bprove\b", problem))
        or bool(re.search(r"\b(?:prove|show)\s+that\b", problem))
    )


def reward_MATH(batch, responses, num_generations, device):
    answers = batch['answers'] * num_generations
    canonical_flags = list(batch.get("answer_is_canonical", [False] * len(batch["answers"]))) * num_generations
    # answer rewards
    ext_ans = [
        normalize_reference_answer(ans, canonical=bool(is_canonical))
        for ans, is_canonical in zip(answers, canonical_flags)
    ]
    # Responses still use the eval-style extractor. Canonical references,
    # however, must not fall back to "last number" because raw answers such as
    # 2\pi - 4 would otherwise become just 4.
    ext_res = [
        extract_answer(res, skip_unit=bool(is_canonical))
        for res, is_canonical in zip(responses, canonical_flags)
    ]
    rewards = torch.zeros(len(answers), device=device)
    for i, (ans, res) in enumerate(zip(ext_ans, ext_res)):
        ans_s = "" if ans is None else str(ans).strip()
        res_s = "" if res is None else str(res).strip()
        if ans_s and res_s and math_equal(ans_s, res_s, timeout=True):
            rewards[i] += 1.0
        else:
            rewards[i] -= 1.0

    return rewards


def load_math_dataset_and_reward(
    local_path: str,
    batch_size: int,
    split: str = 'train', 
    num_workers: int = 8,
    min_level: int = None,
    max_level: int = None,
    only_level: int = None,
    max_rows: int = 1e8,
    rank: int = None,
    num_replicas: int = None,
    seed: int = 112,
):
    ds = load_dataset_split(local_path, split=split)
    # Disable caching to avoid .arrow cache race conditions in multi-GPU
    disable_caching()
    # level <= 2: ~1344
    if min_level is not None:
        ds = ds.filter(lambda x: try_get_level(x['level'], 5) >= min_level)
    if max_level is not None:
        ds = ds.filter(lambda x: try_get_level(x['level'], 5) <= max_level)
    if only_level is not None:
        ds = ds.filter(lambda x: try_get_level(x['level'], 5) == only_level)
    ds = ds.select(range(min(len(ds), max_rows)))
    ds = ds.filter(lambda x: len(x.get('problem', [])) > 0 and len(x.get('problem', '')) < 1500)
    enable_caching()
    ds = ds.with_format('torch')
    ds = ds.shuffle(seed=seed)
    if rank is not None and num_replicas is not None:
        sampler = InfiniteSampler(
            ds, rank=rank, num_replicas=num_replicas, 
        )
    else:
        sampler = InfiniteSampler(
            ds, rank=get_rank(), num_replicas=get_world_size(), 
        )
    
    loader_kwargs = dict(
        collate_fn=collate_fn_math,
        batch_size=batch_size,
        sampler=sampler,
        num_workers=num_workers,
        pin_memory=True,
    )
    if num_workers and int(num_workers) > 0:
        loader_kwargs["persistent_workers"] = True
        loader_kwargs["prefetch_factor"] = 4
    dl = DataLoader(ds, **loader_kwargs)
    
    return dl, reward_MATH


def _load_and_normalize_source(source):
    source = dict(source)
    local_path = source.pop("local_path")
    split = source.pop("split", "train")
    max_rows = int(source.pop("max_rows", 10**9))
    source_name = source.pop("source_dataset", Path(str(local_path)).name)
    require_parseable = bool(source.pop("require_parseable_reward_reference", True))
    max_reference_answer_chars = int(source.pop("max_reference_answer_chars", 256))
    exclude_proof_like = bool(source.pop("exclude_proof_like", False))
    ds = load_dataset_split(local_path, split=split, data_dir=source.pop("data_dir", None))
    if max_rows > 0:
        ds = ds.select(range(min(len(ds), max_rows)))
    ds = ds.map(
        lambda item: normalize_math_record(item, source_name),
        remove_columns=ds.column_names,
        features=ACDIR_REFERENCE_FEATURES,
    )
    ds = ds.filter(lambda x: len(x.get("problem", "")) > 0 and len(x.get("solution", "")) > 0)
    if exclude_proof_like:
        ds = ds.filter(lambda x: not is_proof_like_problem(x))
    if require_parseable:
        ds = ds.filter(lambda x: has_parseable_reward_reference(x, max_chars=max_reference_answer_chars))
    ds = ds.cast(ACDIR_REFERENCE_FEATURES)
    return ds


def _load_profiled_source(source):
    source = dict(source)
    local_path = source.pop("local_path")
    split = source.pop("split", "train")
    max_rows = int(source.pop("max_rows", 10**9))
    exclude_parse_failed = bool(source.pop("exclude_parse_failed", False))
    ds = load_dataset_split(local_path, split=split, data_dir=source.pop("data_dir", None))
    missing = sorted(PROFILE_COLUMNS.difference(ds.column_names))
    if missing:
        raise ValueError(f"Profiled ACDiR dataset is missing columns: {missing}")
    if exclude_parse_failed:
        ds = ds.filter(lambda x: not bool(x.get("parse_failed", False)))
    if max_rows > 0:
        ds = ds.select(range(min(len(ds), max_rows)))
    ds = ds.filter(lambda x: len(x.get("problem", "")) > 0 and len(x.get("solution", "")) > 0)
    return ds


def load_acdir_mixed_reference_dataset_and_reward(
    sources,
    batch_size: int,
    split: str = "train",
    num_workers: int = 8,
    rank: int = None,
    num_replicas: int = None,
    seed: int = 112,
    max_rows: int = 10**9,
):
    del split
    disable_caching()
    datasets = []
    for source in sources:
        datasets.append(_load_and_normalize_source(source))
    if not datasets:
        raise ValueError("sources must be non-empty.")
    ds = concatenate_datasets(datasets) if len(datasets) > 1 else datasets[0]
    if max_rows is not None and int(max_rows) > 0:
        ds = ds.select(range(min(len(ds), int(max_rows))))
    enable_caching()
    ds = ds.shuffle(seed=seed)
    if rank is not None and num_replicas is not None:
        sampler = InfiniteSampler(ds, rank=rank, num_replicas=num_replicas)
    else:
        sampler = InfiniteSampler(ds, rank=get_rank(), num_replicas=get_world_size())
    loader_kwargs = dict(
        collate_fn=collate_fn_acdir_reference,
        batch_size=batch_size,
        sampler=sampler,
        num_workers=num_workers,
        pin_memory=True,
    )
    if num_workers and int(num_workers) > 0:
        loader_kwargs["persistent_workers"] = True
        loader_kwargs["prefetch_factor"] = 4
    return DataLoader(ds, **loader_kwargs), reward_MATH


def load_acdir_profiled_reference_dataset_and_reward(
    sources,
    batch_size: int,
    split: str = "train",
    num_workers: int = 8,
    rank: int = None,
    num_replicas: int = None,
    seed: int = 112,
    max_rows: int = 10**9,
    profile_bucket_weights=None,
    exclude_parse_failed: bool = True,
):
    del split
    disable_caching()
    datasets = []
    for source in sources:
        source = dict(source)
        source.setdefault("exclude_parse_failed", exclude_parse_failed)
        datasets.append(_load_profiled_source(source))
    if not datasets:
        raise ValueError("sources must be non-empty.")
    ds = concatenate_datasets(datasets) if len(datasets) > 1 else datasets[0]
    if exclude_parse_failed:
        ds = ds.filter(lambda x: not bool(x.get("parse_failed", False)))
    if max_rows is not None and int(max_rows) > 0:
        ds = ds.select(range(min(len(ds), int(max_rows))))
    enable_caching()
    if rank is None:
        try:
            rank = get_rank()
        except Exception:
            rank = 0
    else:
        rank = int(rank)
    if num_replicas is None:
        try:
            num_replicas = get_world_size()
        except Exception:
            num_replicas = 1
    else:
        num_replicas = int(num_replicas)
    if profile_bucket_weights:
        sampler = ProfileBucketSampler(
            ds,
            profile_bucket_weights,
            rank=rank,
            num_replicas=num_replicas,
            seed=seed,
            exclude_parse_failed=exclude_parse_failed,
        )
    else:
        ds = ds.shuffle(seed=seed)
        sampler = InfiniteSampler(ds, rank=rank, num_replicas=num_replicas)
    loader_kwargs = dict(
        collate_fn=collate_fn_acdir_profiled_reference,
        batch_size=batch_size,
        sampler=sampler,
        num_workers=num_workers,
        pin_memory=True,
    )
    if num_workers and int(num_workers) > 0:
        loader_kwargs["persistent_workers"] = True
        loader_kwargs["prefetch_factor"] = 4
    return DataLoader(ds, **loader_kwargs), reward_MATH


def extract_answer_gsm8k(answer: str):
    # find the last part starting with '#### xxx'
    return answer.split('####')[-1].strip()


def reward_gsm8k(
    batch, responses, num_generations, device
):
    answers = batch['answers'] * num_generations
    # answer rewards
    ext_ans = [extract_answer_gsm8k(ans) for ans in answers]
    ext_res = [extract_answer(res) for res in responses]
    rewards = torch.zeros(len(answers), device=device)
    for i, (ans, res) in enumerate(zip(ext_ans, ext_res)):
        if math_equal(ans, res):
            rewards[i] += 1.0
        else:
            rewards[i] -= 1.0

    return rewards

def load_gsm8k_dataset_and_reward(
    local_path: str,
    batch_size: int,
    split: str = 'train', 
    num_workers: int = 8,
    rank: int = None,
    num_replicas: int = None,
    seed: int = 112, 
):
    ds = load_dataset_split(local_path, split=split, data_dir='main')
    ds = ds.with_format('torch')
    ds = ds.shuffle(seed=seed)
    if rank is not None and num_replicas is not None:
        sampler = InfiniteSampler(
            ds, rank=rank, num_replicas=num_replicas, 
        )
    else:
        sampler = InfiniteSampler(
            ds, rank=get_rank(), num_replicas=get_world_size(), 
        )

    loader_kwargs = dict(
        collate_fn=collate_fn_gsm8k,
        batch_size=batch_size,
        sampler=sampler,
        num_workers=num_workers,
        pin_memory=True,
    )
    if num_workers and int(num_workers) > 0:
        loader_kwargs["persistent_workers"] = True
        loader_kwargs["prefetch_factor"] = 4
    dl = DataLoader(ds, **loader_kwargs)

    return dl, reward_gsm8k