diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index b13e6f9c78..8f60fa3ca9 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -613,17 +613,17 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import
 
     # Remove non editted functions
     for remove in remover:
-        RLTrainer_source = RLTrainer_source.replace(remove, "\n")
+        RLTrainer_source = RLTrainer_source.replace(remove, "")
     pass
-    
+
     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}:\n    def __init__(",
+        rf"class _Unsloth{RLTrainer_name}(.*?:).+?def __init__\(",
+        rf"class _Unsloth{RLTrainer_name}\1\n    def __init__(",
         RLTrainer_source,
         flags = re.MULTILINE | re.DOTALL,
     )