diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 2342574737..d70a83f71d 100644 --- a/unsloth/models/rl.py +++ b/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