From 19a0c3fe3205837e9a0879b144fbb76920feeebe Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 13 Feb 2025 04:25:56 -0800
Subject: [PATCH] Update rl_replacements.py

---
 unsloth/models/rl_replacements.py | 24 +++++++++++++-----------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py
index 3eb16bb1f4..968f2b19fd 100644
--- a/unsloth/models/rl_replacements.py
+++ b/unsloth/models/rl_replacements.py
@@ -18,6 +18,7 @@ __all__ = [
 ]
 
 import re
+import inspect
 from collections import defaultdict
 RL_EXTRA_ARGS = defaultdict(list)
 RL_FUNCTIONS  = defaultdict(list)
@@ -99,20 +100,21 @@ RL_FUNCTIONS["sft_trainer"].append(sft_trainer_prepare_dataset)
 
 
 # Ignore mean_token_accuracy since it needs logits
+# We override it directly with our version
+def _sft_trainer_compute_loss(self, model, inputs, return_outputs = False, num_items_in_batch = None):
+    (loss, outputs) = super().compute_loss(
+        model,
+        inputs,
+        return_outputs = return_outputs,
+        num_items_in_batch = num_items_in_batch,
+    )
+    return (loss, outputs) if return_outputs else loss
+pass
+
 def sft_trainer_compute_loss(function_name, function):
     if  function_name != "compute_loss": return function
 
-    # .*? matches first match. .+? matches final match.
-    replacer = re.findall(
-        r"\.compute_loss\(.*?\)",
-        function,
-        flags = re.MULTILINE | re.DOTALL,
-    )
-    if len(replacer) != 0:
-        replacer = replacer[0]
-        returner = "\n" + " "*8 + "return (loss, outputs) if return_outputs else loss"
-        function = function.replace(replacer, replacer + returner)
-    pass
+    function = inspect.getsource(_sft_trainer_compute_loss)
     return function
 pass
 RL_FUNCTIONS["sft_trainer"].append(sft_trainer_compute_loss)