temp

Author: A9isha

src/maxtext/trainers/post_train/rl/train_rl.py

@@ -121,29 +121,18 @@ def get_dataset(
   if dataset_name is None:
     raise ValueError("dataset_name must be provided")
 
-  if dataset_name.startswith("huggingface:"):
-    import datasets  # pylint: disable=import-outside-toplevel
-
-    if data_files is None:
-      hf_dataset_name = dataset_name.replace("huggingface:", "")
-      data = datasets.load_dataset(hf_dataset_name, split=split, cache_dir=data_dir)
-      if tmvp_config.debug.rl:
-        max_logging.log(f"Loaded Hugging Face dataset {hf_dataset_name} with split {split}. Size: {len(data)}")
-    else:  # data_files have been provided, useful for using slices of large datasets like nvidia/OpenMathInstruct-2
-      data = datasets.load_dataset(
-          "parquet",
-          data_files={tmvp_config.train_split: data_files},
-          split=split,
-          cache_dir=data_dir,
-      )
-  else:
-    builder_kwargs = {"file_format": tfds.core.FileFormat.ARRAY_RECORD}
-    data = tfds.data_source(
-        dataset_name,
+  import datasets  # pylint: disable=import-outside-toplevel
+
+  if data_files is None:
+    data = datasets.load_dataset(dataset_name, split=split, cache_dir=data_dir)
+    if tmvp_config.debug.rl:
+      max_logging.log(f"Loaded Hugging Face dataset {dataset_name} with split {split}. Size: {len(data)}")
+  else:  # data_files have been provided, useful for using slices of large datasets like nvidia/OpenMathInstruct-2
+    data = datasets.load_dataset(
+        "parquet",
+        data_files={tmvp_config.train_split: data_files},
         split=split,
-        data_dir=data_dir,
-        builder_kwargs=builder_kwargs,
-        download=True,
+        cache_dir=data_dir,
     )
 
   template_config = load_template_from_file(tmvp_config.chat_template_path)
@@ -300,6 +289,37 @@ def get_max_train_steps(trainer_config):
   )
 
 
+def prepare_train_and_eval_dataset(
+    trainer_config,
+    seed: int = 42,
+    test_size: float = 0.05,
+):
+  """Load and split the dataset into train and validation sets using HF's train_test_split."""
+  import datasets  # pylint: disable=import-outside-toplevel
+
+  max_logging.log(
+      "WARNING: For reproducible experiments, preprocess the dataset once and "
+      "define your own HfDataset subclass that directly uses the preprocessed datasets."
+  )
+
+  original_ds = datasets.load_dataset(
+      "parquet",
+      data_files={trainer_config.train_split: trainer_config.hf_train_files},
+      split=trainer_config.train_split,
+  )
+
+  if "OpenMathReasoning" in trainer_config.dataset_name:
+    original_ds = original_ds.filter(lambda x: x.get("problem_type") == "has_answer_extracted")
+
+  # Split into train and validation sets using HF's train_test_split
+  split_ds = original_ds.train_test_split(test_size=test_size, seed=seed)
+
+  return {
+      "train": split_ds["train"],
+      "validation": split_ds["test"],
+  }
+
+
 def prepare_datasets(trainer_config, model_tokenizer):
   """Setup and return train and test datasets."""
   home = os.path.expanduser("~") + "/"
@@ -311,39 +331,16 @@ def prepare_datasets(trainer_config, model_tokenizer):
     os.makedirs(test_data_dir)
 
   # Prepare train and test data from training data for certain datasets
-  if trainer_config.dataset_name in ["nvidia/OpenMathInstruct-2", "nvidia/OpenMathReasoning", "open-r1/OpenR1-Math-220k", "bethgelab/CuratedThoughts"]:
+  eval_dataset_name = getattr(trainer_config, "eval_dataset_name", None)
+  if trainer_config.dataset_name in [
+      "nvidia/OpenMathInstruct-2",
+      "nvidia/OpenMathReasoning",
+      "open-r1/OpenR1-Math-220k",
+      "bethgelab/CuratedThoughts",
+  ] and (not eval_dataset_name or eval_dataset_name == trainer_config.dataset_name):
     import datasets  # pylint: disable=import-outside-toplevel
 
-    def prepare_train_and_eval_dataset(
-        seed: int = 42,
-        test_size: float = 0.05,
-    ):
-      """Load and split the dataset into train and validation sets using HF's train_test_split."""
-      max_logging.log(
-          "WARNING: For reproducible experiments, preprocess the dataset once and "
-          "define your own HfDataset subclass that directly uses the preprocessed datasets."
-      )
-
-      # Load the original dataset
-      original_ds = datasets.load_dataset(
-          "parquet",
-          data_files={trainer_config.train_split: trainer_config.hf_train_files},
-          split=trainer_config.train_split,
-      )
-
-      if "OpenMathReasoning" in trainer_config.dataset_name:
-            original_ds = original_ds.filter(lambda x: x.get("problem_type") == "has_answer_extracted")
-
-
-      # Split into train and validation sets using HF's train_test_split
-      split_ds = original_ds.train_test_split(test_size=test_size, seed=seed)
-
-      return {
-          "train": split_ds["train"],
-          "validation": split_ds["test"],
-      }
-
-    splits = prepare_train_and_eval_dataset()
+    splits = prepare_train_and_eval_dataset(trainer_config)
     template_config = load_template_from_file(trainer_config.chat_template_path)
 
     train_dataset = (

src/maxtext/trainers/post_train/rl/utils_rl.py

@@ -42,10 +42,13 @@
     LatexExtractionConfig(),
 )
 
+
 def math_verify_func(items, timeout=5):
   """Verifies a batch of math problems, handling timeouts and exceptions."""
   with concurrent.futures.ThreadPoolExecutor() as executor:
-    future_to_index = {executor.submit(verify_math, golds, predictions): idx for idx, (_, golds, predictions) in enumerate(items)}
+    future_to_index = {
+        executor.submit(verify_math, golds, predictions): idx for idx, (_, golds, predictions) in enumerate(items)
+    }
     results = [0.0] * len(items)
     for future in concurrent.futures.as_completed(future_to_index):
       index = future_to_index[future]
@@ -59,8 +62,12 @@ def math_verify_func(items, timeout=5):
 def verify_math(golds, predictions):
   """Runs mathematical expression evaluation on ground-truth and predictions."""
 
-  extracted_predictions = list(itertools.chain.from_iterable(parse(pred, PRED_EXTRACTION_TARGET, parsing_timeout=None) for pred in predictions))
-  extracted_golds = list(itertools.chain.from_iterable(parse(gold, GOLD_EXTRACTION_TARGET, parsing_timeout=None) for gold in golds))
+  extracted_predictions = list(
+      itertools.chain.from_iterable(parse(pred, PRED_EXTRACTION_TARGET, parsing_timeout=None) for pred in predictions)
+  )
+  extracted_golds = list(
+      itertools.chain.from_iterable(parse(gold, GOLD_EXTRACTION_TARGET, parsing_timeout=None) for gold in golds)
+  )
   # If no predictions or golds were extracted, return 0.0
   if not extracted_predictions or not extracted_golds:
     return 0.0
@@ -72,6 +79,7 @@ def verify_math(golds, predictions):
       ]
   )
 
+
 def boxed(x):
   """Wraps the input string in a LaTeX boxed command if it's not already wrapped."""
   return "\\boxed{" + x + "}" if not x.startswith("\\boxed{") else x
@@ -267,7 +275,10 @@ def normalize_final_answer(final_answer: str) -> str:
 def preprocess_math_string(dataset_name, text) -> str:
   """Fix common formatting issues in text."""
   # Normalize for certain datasets and parse
-  if any(name in dataset_name for name in ["DAPO", "OpenMathInstruct", "OpenMathReasoning",  "OpenR1-Math-220k", "CuratedThoughts"]):
+  if any(
+      name in dataset_name
+      for name in ["DAPO", "OpenMathInstruct", "OpenMathReasoning", "OpenR1-Math-220k", "CuratedThoughts", "MATH-500"]
+  ):
     text = normalize_final_answer(text).strip()
   # Fix LaTeX escaping issues
   text = fix_latex_escaping(text)
@@ -418,7 +429,11 @@ def check_numbers(prompts, completions, answer, tmvp_config, **kargs):
         # 3. As a fallback, try numeric comparison if both can be parsed as numbers
         try:
           predictions = parse(norm_guesses[0], PRED_EXTRACTION_TARGET, parsing_timeout=None)
-          golds = list(itertools.chain.from_iterable(parse(norm_answer, GOLD_EXTRACTION_TARGET, parsing_timeout=None) for norm_answer in norm_answers))
+          golds = list(
+              itertools.chain.from_iterable(
+                  parse(norm_answer, GOLD_EXTRACTION_TARGET, parsing_timeout=None) for norm_answer in norm_answers
+              )
+          )
           for gold in golds:
             for pred in predictions:
               try:
@@ -430,9 +445,11 @@ def check_numbers(prompts, completions, answer, tmvp_config, **kargs):
                 else:
                   scores[gen_idx] = max(scores[gen_idx], tmvp_config.penalty_incorrect_answer)
               except:
-                  scores[gen_idx] = max(scores[gen_idx], tmvp_config.penalty_incorrect_answer)
+                scores[gen_idx] = max(scores[gen_idx], tmvp_config.penalty_incorrect_answer)
         except:
-          scores[gen_idx] = max(scores[gen_idx], tmvp_config.penalty_incorrect_format)  # Penalize if we can't parse numbers at all
+          scores[gen_idx] = max(
+              scores[gen_idx], tmvp_config.penalty_incorrect_format
+          )  # Penalize if we can't parse numbers at all
   if tmvp_config.debug.rl:
     debug_log_path = epath.Path(tmvp_config.base_output_directory) / tmvp_config.run_name / "debug_rl_logs"
     debug_log_path.mkdir(parents=True, exist_ok=True)
@@ -469,10 +486,11 @@ def extract_hash_answer(text: str) -> str | None:
 def check_correctness(extracted_response, acceptable_answers, tmvp_config):
   """Handles math verification and partial correctness logic."""
   norm_answers = []
-  norm_response = preprocess_math_string(tmvp_config.dataset_name, extracted_response)
+  dataset_name = tmvp_config.eval_dataset_name if tmvp_config.eval_dataset_name else tmvp_config.dataset_name
+  norm_response = preprocess_math_string(dataset_name, extracted_response)
   # Check exact correctness first
-  for answer in acceptable_answers: 
-    norm_answers.append(preprocess_math_string(tmvp_config.dataset_name, answer))
+  for answer in acceptable_answers:
+    norm_answers.append(preprocess_math_string(dataset_name, answer))
   is_correct = verify_math([boxed(norm_answer) for norm_answer in norm_answers], [boxed(norm_response)]) > 0.1
   if is_correct:
     return True, True  # Exact correctness implies partial correctness
@@ -481,7 +499,11 @@ def check_correctness(extracted_response, acceptable_answers, tmvp_config):
   is_partially_correct = False
   try:
     predictions = parse(boxed(norm_response), PRED_EXTRACTION_TARGET, parsing_timeout=None)
-    golds = list(itertools.chain.from_iterable(parse(boxed(norm_answer), GOLD_EXTRACTION_TARGET, parsing_timeout=None) for norm_answer in norm_answers))
+    golds = list(
+        itertools.chain.from_iterable(
+            parse(boxed(norm_answer), GOLD_EXTRACTION_TARGET, parsing_timeout=None) for norm_answer in norm_answers
+        )
+    )
     is_partially_correct = any(
         0.9 <= (float(pred) + EPSILON) / (float(gold) + EPSILON) <= 1.1 for pred in predictions for gold in golds
     )