From ecaf62887f7ae3dea60ecb775b971d4334624917 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sat, 15 Feb 2025 17:57:47 -0800
Subject: [PATCH] Update rl.py

---
 unsloth/models/rl.py | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index 8f60fa3ca9..51a5abb75d 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -457,6 +457,11 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
         selective_log_softmax_code = selective_log_softmax_code,
     )
 
+    # Remove multiple doc strings
+    if RLTrainer_source.count(__RLTrainer_doc__) == 2:
+        RLTrainer_source = RLTrainer_source.replace(__RLTrainer_doc__, "", 1)
+    pass
+
     # Create new function
     created_module = create_new_function(
         f"Unsloth{RLTrainer_name}",
@@ -619,14 +624,6 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import
     RLTrainer_source = RLTrainer_source.replace(
         f"class {RLTrainer_name}", f"class _Unsloth{RLTrainer_name}", 1
     )
-
-    # Get rid of docs since we repeated it
-    RLTrainer_source = re.sub(
-        rf"class _Unsloth{RLTrainer_name}(.*?:).+?def __init__\(",
-        rf"class _Unsloth{RLTrainer_name}\1\n    def __init__(",
-        RLTrainer_source,
-        flags = re.MULTILINE | re.DOTALL,
-    )
     return RLTrainer_source
 pass