Merge pull request #3324 from AI-Hypercomputer:hengtaoguo-parse

Google-ML-Automation · Google-ML-Automation · commit 146fb8ec53ae · 2026-03-05T17:53:02.000-08:00
PiperOrigin-RevId: 879335972
diff --git a/src/maxtext/configs/post_train/rl.yml b/src/maxtext/configs/post_train/rl.yml
@@ -153,7 +153,7 @@ max_num_seqs: null
 # If True, enables asynchronous scheduling in vLLM for faster generation
 async_scheduling: True
 # stop generation when any of these strings is generated
-stop_strings: [</answer>]
+stop_strings: null
 
 # ====== Checkpoint Configuration ======
 enable_checkpointing: True
diff --git a/src/maxtext/trainers/post_train/rl/utils_rl.py b/src/maxtext/trainers/post_train/rl/utils_rl.py
@@ -106,7 +106,7 @@ def get_match_format_regex(tmvp_config):
   match_format = re.compile(
       (
           r"^[\s]{0,}"
-          rf"{tmvp_config.reasoning_start_token}.+?{tmvp_config.reasoning_end_token}.*?"
+          rf"{tmvp_config.reasoning_start_token}.+{tmvp_config.reasoning_end_token}.*?"
           rf"{tmvp_config.solution_start_token}(.+?){tmvp_config.solution_end_token}"
       ),
       flags=re.MULTILINE | re.DOTALL,

Original file line number	Diff line number	Diff line change
`@@ -106,7 +106,7 @@ def get_match_format_regex(tmvp_config):`
`106`	`106`	`match_format = re.compile(`
`107`	`107`	`(`
`108`	`108`	`r"^[\s]{0,}"`
`109`		`- rf"{tmvp_config.reasoning_start_token}.+?{tmvp_config.reasoning_end_token}.*?"`
	`109`	`+ rf"{tmvp_config.reasoning_start_token}.+{tmvp_config.reasoning_end_token}.*?"`
`110`	`110`	`rf"{tmvp_config.solution_start_token}(.+?){tmvp_config.solution_end_token}"`
`111`	`111`	`),`
`112`	`112`	`flags=re.MULTILINE \| re.DOTALL,`