Extra replacements

unsloth/models/rl.py

@@ -23,7 +23,9 @@ import os
 import re
 from unsloth_zoo.compiler import create_new_function
 from unsloth_zoo.logging_utils import PatchRLStatistics
-
+from .rl_replacements import (
+    RL_EXTRA_ARGS,
+)
 
 def PatchRL(FastLanguageModel):
 
@@ -282,6 +284,13 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
         "from unsloth_zoo.logging_utils import PatchRLStatistics\n"\
         f"PatchRLStatistics('{trainer_file}')\n"
 
+    # Patch optional args
+    if trainer_file in RL_EXTRA_ARGS:
+        process_extra_args = RL_EXTRA_ARGS[trainer_file]
+        for process_extra_arg in process_extra_args:
+            extra_args += process_extra_args(call_args, extra_args)
+    pass
+
     # Create RLTrainer args
     extra_args = extra_args.split("\n")
     extra_args = "\n".join(" "*8 + x for x in extra_args)

unsloth/models/rl_replacements.py

@@ -0,0 +1,50 @@
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+__all__ = [
+    "RL_EXTRA_ARGS",
+]
+
+RL_EXTRA_ARGS = dict()
+
+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"\
+        "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"\
+        "fix_untrained_tokens(model, tokenizer, train_dataset, IGNORED_TOKENIZER_NAMES, eps = 1e-16)\n"\
+        "fix_zero_training_loss(model, tokenizer, train_dataset)\n"
+        return fix_tokenizer
+    return ""
+pass
+RL_EXTRA_ARGS["sft_trainer"] = [sft_trainer_fix_untraiend_tokens,]
+
+
+def dpo_trainer_fix_columns(call_args, extra_args):
+    if "model" in call_args and "train_dataset" in call_args:
+        fix_dpo = \
+        "if hasattr(train_dataset, 'column_names'):\n"\
+        "    column_names = set(train_dataset.column_names)\n"\
+        "    check = ['chosen', 'rejected', 'prompt', 'chosen_input_ids', 'chosen_attention_mask',\n"\
+        "             'chosen_labels', 'rejected_input_ids', 'rejected_attention_mask', 'rejected_labels',\n"\
+        "             'prompt_input_ids', 'prompt_attention_mask']\n"\
+        "    if all(x in column_names for x in check):\n"\
+        "        train_dataset = train_dataset.remove_columns(['chosen', 'rejected', 'prompt'])\n"\
+        "    del check, column_names\n"\
+        return fix_dpo
+    return ""
+pass
+RL_EXTRA_ARGS["dpo_trainer"] = [dpo_trainer_fix_columns,]

unsloth/tokenizer_utils.py

@@ -59,6 +59,7 @@ IGNORED_TOKENIZER_NAMES = frozenset(
     [x.lower() for x in IGNORED_TOKENIZER_NAMES] + \
     [x.lower()+"-bnb-4bit" for x in IGNORED_TOKENIZER_NAMES]
 )
+os.environ["UNSLOTH_IGNORED_TOKENIZER_NAMES"] = "\n".join(IGNORED_TOKENIZER_NAMES)
 
 # Check environments
 keynames = "\n" + "\n".join(os.environ.keys())
@@ -1055,5 +1056,5 @@ def patch_sft_trainer_tokenizer():
     pass
 pass
 
-# Finally patch TRL tokenizer things
+# Finally patch TRL tokenizer things -> moved to RL
 # patch_sft_trainer_tokenizer()