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