Add open-r1/OpenR1-Math-220k dataset to RL

Author: SurbhiJainUSC

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

@@ -100,8 +100,6 @@ def score_responses(tmvp_config, question, responses, answers):
   Returns:
       Tuple of (is_correct, is_partially_correct, has_correct_format)
   """
-
-  answers = list(dict.fromkeys(answers))
   if tmvp_config.debug.rl:
     max_logging.log("========================================")
     max_logging.log(f"Evaluation Question: {question}")
@@ -174,17 +172,14 @@ def evaluate(
 
     # Generate responses for all prompts in the batch
     multiple_call_responses = generate_responses(
-        tmvp_config=tmvp_config,
         prompts=prompts,
         rl_cluster=rl_cluster,
         num_passes=num_passes,
     )
 
     # Score each question-answer pair
     for question, responses, answer in zip(questions, multiple_call_responses, answers):
-      answer = (
-          json.loads(answer) if isinstance(answer, str) else answer
-      )  # decode the json-encoded list of acceptable answers
+      answer = list(dict.fromkeys(json.loads(answer)))
       is_correct, is_partially_correct, has_correct_format = score_responses(
           tmvp_config=tmvp_config,
           question=question,

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

@@ -324,7 +324,7 @@ def prepare_train_and_eval_dataset(
       )
 
       original_ds = datasets.load_dataset(
-          trainer_config.dataset_name,
+          "parquet",
           data_files={trainer_config.train_split: trainer_config.hf_train_files},
           split=trainer_config.train_split,
       )

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

@@ -379,8 +379,8 @@ def check_numbers(prompts, completions, answer, tmvp_config, **kargs):
 
   # Extract full answer content from solution tags (not just first number)
   extracted_responses = [extract_answer(c, tmvp_config) for c in completions]
-  true_answers = [list(dict.fromkeys(json.loads(acceptable_answers))) if isinstance(acceptable_answers, str) else acceptable_answers for acceptable_answers in answer]
-
+  true_answers = [list(dict.fromkeys(acceptable_answers)) for acceptable_answers in answer]
+  
   if tmvp_config.debug.rl:
     max_logging.log("START ============================")
     max_logging.log(f"Question: {question[0]}")
@@ -417,24 +417,26 @@ def check_numbers(prompts, completions, answer, tmvp_config, **kargs):
   if math_verify_queue:
     # 2. Try math_verify for robust mathematical correctness checking
     math_verify_results = math_verify_func(math_verify_queue)
-    for (gen_idx, norm_answers, norm_guess), score in zip(math_verify_queue, math_verify_results):
+    for (gen_idx, norm_answers, norm_guesses), score in zip(math_verify_queue, math_verify_results):
       if score > 0.1:
         scores[gen_idx] = max(scores[gen_idx], tmvp_config.reward_exact_answer)
-        print("-------- Found a math_verify match -----------")
       else:
         # 3. As a fallback, try numeric comparison if both can be parsed as numbers
         try:
-          predictions = parse(boxed(norm_guess), 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))
+          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))
           for gold in golds:
             for pred in predictions:
-              ratio = (float(pred) + EPSILON) / (float(gold) + EPSILON)
-              if 0.9 <= ratio <= 1.1:
-                scores[gen_idx] = max(scores[gen_idx], tmvp_config.reward_ratio_guess_to_answer_high)
-              elif 0.8 <= ratio <= 1.2:
-                scores[gen_idx] = max(scores[gen_idx], tmvp_config.reward_ratio_guess_to_answer_low)
-              else:
-                scores[gen_idx] = max(scores[gen_idx], tmvp_config.penalty_incorrect_answer)  # Penalize wrong answers
+              try:
+                ratio = (float(pred) + EPSILON) / (float(gold) + EPSILON)
+                if 0.9 <= ratio <= 1.1:
+                  scores[gen_idx] = max(scores[gen_idx], tmvp_config.reward_ratio_guess_to_answer_high)
+                elif 0.8 <= ratio <= 1.2:
+                  scores[gen_idx] = max(scores[gen_idx], tmvp_config.reward_ratio_guess_to_answer_low)
+                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)
         except:
           scores[gen_idx] = max(scores[gen_idx], tmvp_config.penalty_incorrect_format)  # Penalize if we can't parse numbers at all
   return scores
@@ -462,20 +464,22 @@ def check_correctness(extracted_response, acceptable_answers, tmvp_config):
   # Check exact correctness first
   for answer in acceptable_answers: 
     norm_answers.append(preprocess_math_string(tmvp_config.dataset_name, answer))
-  is_correct = math_verify_func([(0, [boxed(norm_answer) for norm_answer in norm_answers], [boxed(norm_response)])])[0] > 0.1
+  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
 
   # Check partial correctness if values can be extracted (within 10%)
-  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))
-  is_partially_correct = any(
-      0.9 <= (float(pred) + EPSILON) / (float(gold) + EPSILON) <= 1.1 for pred in predictions for gold in golds
-  )
-  if is_partially_correct:
-    return False, True  # Not exactly correct, but partially correct
+  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))
+    is_partially_correct = any(
+        0.9 <= (float(pred) + EPSILON) / (float(gold) + EPSILON) <= 1.1 for pred in predictions for gold in golds
+    )
+  except:
+    pass
 
-  return False, False  # Not correct at all
+  return False, is_partially_correct
 
 
 def get_optimizer(tmvp_config, max_train_steps):