Update rl.py

unsloth/models/rl.py

@@ -38,6 +38,7 @@ import inspect
 import os
 import re
 import functools
+from unsloth_zoo.compiler import create_new_function
 
 
 def PatchRL(FastLanguageModel):
@@ -319,9 +320,6 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
     __init__ = __init__.replace("if peft_config is not None:", "if False:")
     __init__ = __init__.replace("get_peft_model(model, peft_config)", "model")
 
-    # Change super() to Trainer
-    __init__ = __init__.replace("super()", f"super(Unsloth{RLTrainer_name}, self)")
-
     # Add spaces back into __init__
     __init__ = __init__.split("\n")
     __init__ = "\n".join(' '*spaces + x for x in __init__)
@@ -374,9 +372,6 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
 
     # Import all functions
     imports = list(set(imports))
-    imports = f"from trl.trainer.{trainer_file} import (\n" + ',\n'.join(imports) + ")"
-    imported_functions = {}
-    exec(imports, globals(), imported_functions)
 
     # Patch all functions
     for function in changed:
@@ -386,11 +381,14 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
     RLTrainer_source = RLTrainer_source.replace(
         f"class {RLTrainer_name}", f"class Unsloth{RLTrainer_name}", 1
     )
-    exec(RLTrainer_source, imported_functions, globals())
-    globals()[f"Unsloth{RLTrainer_name}"] = eval(f"Unsloth{RLTrainer_name}")
-    exec(f"trl.trainer.{trainer_file}.{RLTrainer_name} = Unsloth{RLTrainer_name}", locals(), globals())
-    exec(f"trl.trainer.{RLTrainer_name} = Unsloth{RLTrainer_name}", locals(), globals())
-    exec(f"trl.{RLTrainer_name} = Unsloth{RLTrainer_name}", locals(), globals())
+
+    module = create_new_function(
+        RLTrainer_name,
+        RLTrainer_source,
+        f"trl.trainer.{trainer_file}",
+        imports,
+    )
+    return module
 pass