Update rl.py

This commit is contained in:
Daniel Han 2025-02-15 17:48:52 -08:00
commit 742450a741

View file

@ -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,
)