From fe21809bda0fed9064081094a057598aea7b78ca Mon Sep 17 00:00:00 2001 From: numb3r33 Date: Wed, 24 Dec 2025 00:14:50 +0530 Subject: [PATCH] Fix indentation handling in grpo_trainer return statement replacement Use regex to dynamically detect and preserve the original indentation when replacing the 'return output' statement, instead of hardcoding spaces. This ensures the patched code maintains consistent indentation regardless of the original formatting. --- unsloth/models/rl_replacements.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index e13e5e6d78..bcac699b3f 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -388,13 +388,17 @@ def grpo_trainer__generate_and_score_completions(function_name, function): patched = patched[: match.start()] + wrapped + patched[match.end() :] function = patched + + match = re.search(r'^(\s*)return output', function, re.MULTILINE) - function = function.replace( - " return output", - """ if not _was_training: - self.model.for_inference() - return output""", - ) + if match: + indent = match.group(1) + function = function.replace( + f"{indent}return output", + f"""{indent}if not _was_training: + {indent} self.model.for_inference() + {indent}return output""" + ) return function