From 89f2d3a28b0744edc5476eb2341ee91bd4621a81 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