extra RL replacements

unsloth/models/rl.py

@@ -25,6 +25,7 @@ from unsloth_zoo.compiler import create_new_function
 from unsloth_zoo.logging_utils import PatchRLStatistics
 from .rl_replacements import (
     RL_EXTRA_ARGS,
+    RL_FUNCTIONS,
 )
 
 def PatchRL(FastLanguageModel):
@@ -365,8 +366,8 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
     RLConfig_extra_args = extra_args
     RLConfig_call_args  = call_args
 
-    # Patch vLLM
-    RLTrainer_extras = patch_vllm(RLTrainer, trainer_file, RLTrainer_name, all_imports, imports)
+    # Patch vLLM and other functions
+    RLTrainer_extras = patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, imports)
     if RLTrainer_extras is None:
         RLTrainer_extras = f"_Unsloth{RLTrainer_name} = {RLTrainer_name}"
 
@@ -414,7 +415,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
 pass
 
 
-def patch_vllm(RLTrainer, trainer_file, RLTrainer_name, all_imports, imports):
+def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, imports):
     init = inspect.getsource(RLTrainer.__init__)
     old_init = init
 
@@ -475,6 +476,7 @@ def patch_vllm(RLTrainer, trainer_file, RLTrainer_name, all_imports, imports):
     functions = [x for x in functions if f"def {x}" in RLTrainer_source]
 
     changed = {"__init__" : (old_init, init,)}
+    edit_functions = RL_FUNCTIONS.get(trainer_file, [])
 
     for function in functions:
         if not hasattr(RLTrainer, function): continue
@@ -483,6 +485,11 @@ def patch_vllm(RLTrainer, trainer_file, RLTrainer_name, all_imports, imports):
         except: continue
         original_source = source
 
+        # Check for function
+        for edit_function in edit_functions:
+            source = edit_function(function, source)
+        pass
+
         # llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model
         source = re.sub(
             r"(\n[\s]{4,}).+?model_executor\.driver_worker.+?\n",

unsloth/models/rl_replacements.py

@@ -14,14 +14,19 @@
 
 __all__ = [
     "RL_EXTRA_ARGS",
+    "RL_FUNCTIONS",
 ]
 
-RL_EXTRA_ARGS = dict()
+import re
+from collections import defaultdict
+RL_EXTRA_ARGS = defaultdict(list)
+RL_FUNCTIONS  = defaultdict(list)
+
 
 def sft_trainer_fix_untraiend_tokens(call_args, extra_args):
     if "model" in call_args and "train_dataset" in call_args:
         fix_tokenizer = \
-        "IGNORED_TOKENIZER_NAMES = os.environ.get('UNSLOTH_IGNORED_TOKENIZER_NAMES', set())\n"\
+        "IGNORED_TOKENIZER_NAMES = os.environ.get('UNSLOTH_IGNORED_TOKENIZER_NAMES', '').split('\n')\n"\
         "from unsloth_zoo.tokenizer_utils import fix_untrained_tokens\n"\
         "from unsloth_zoo.training_utils  import fix_zero_training_loss\n"\
         "if 'tokenizer' not in locals(): tokenizer = processing_class\n"\
@@ -30,7 +35,7 @@ def sft_trainer_fix_untraiend_tokens(call_args, extra_args):
         return fix_tokenizer
     return ""
 pass
-RL_EXTRA_ARGS["sft_trainer"] = [sft_trainer_fix_untraiend_tokens,]
+RL_EXTRA_ARGS["sft_trainer"].append(sft_trainer_fix_untraiend_tokens)
 
 
 def dpo_trainer_fix_columns(call_args, extra_args):
@@ -47,4 +52,45 @@ def dpo_trainer_fix_columns(call_args, extra_args):
         return fix_dpo
     return ""
 pass
-RL_EXTRA_ARGS["dpo_trainer"] = [dpo_trainer_fix_columns,]
+RL_EXTRA_ARGS["dpo_trainer"].append(dpo_trainer_fix_columns)
+
+
+def sft_trainer_prepare_dataset(function_name, function):
+    if  function_name != "_prepare_non_packed_dataloader" and \
+        function_name != "_prepare_dataset": return
+    
+    check_text = \
+    "\n"\
+    "if 'tokenizer'          not in locals(): tokenizer = processing_class\n"\
+    "if 'formatting_func'    not in locals(): raise RuntimeError('Unsloth: Please file a bug report - `formatting_func` does not exist!')\n"\
+    "if 'dataset_text_field' not in locals() and 'args' in locals(): dataset_text_field = args.dataset_text_field\n"\
+    "if 'dataset_text_field' not in locals(): raise RuntimeError('Unsloth: Please file a bug report - `dataset_text_field` does not exist!')\n"\
+    "test_text = dataset[0][dataset_text_field] if (formatting_func is None and dataset_text_field is not None) else formatting_func(dataset[0])[0]\n"\
+    "chat_template = getattr(tokenizer, 'chat_template', None)\n"\
+    "chat_template = '' if chat_template is None else chat_template\n"\
+    "has_bos_token_already = (test_text.startswith(tokenizer.bos_token) or tokenizer.bos_token in chat_template) "\
+    "if getattr(tokenizer, 'bos_token', None) is not None else False\n"\
+    "if 'add_special_tokens' not in locals() and has_bos_token_already:\n"\
+    "    from functools import partial\n"\
+    "    tokenizer = partial(tokenizer, add_special_tokens = False)\n"\
+    "    processing_class = tokenizer\n"\
+    "else:\n"\
+    "    add_special_tokens = False if has_bos_token_already else add_special_tokens\n"
+
+    check_text = check_text.split("\n")
+    check_text = "\n".join(" "*where + x for x in check_text)
+    check_text = check_text.rstrip() + "\n"
+
+    # .*? matches first match. .+? matches final match.
+    replacer = re.findall(
+        f"def {function_name}\(.*?\).*?\:\n",
+        function,
+        flags = re.MULTILINE | re.DOTALL,
+    )
+    if len(replacer) != 0:
+        replacer = replacer[0]
+        function = function.replace(replacer, replacer + check_text)
+        pass
+    return function
+pass
+RL_FUNCTIONS["sft_trainer"].append(sft_trainer_prepare_dataset)