From 742450a741bf847ec94153af5f7e796619e4f3a6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 15 Feb 2025 17:48:52 -0800 Subject: [PATCH] Update rl.py --- unsloth/models/rl.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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, )