From 56bb441cf4373e405e06b279103aee2acd4ec6a4 Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Tue, 26 Mar 2024 03:58:26 +1100
Subject: [PATCH] patch_tokenizer

---
 unsloth/models/_utils.py   | 25 +++++++++++++++++++++++++
 unsloth/tokenizer_utils.py | 29 ++---------------------------
 2 files changed, 27 insertions(+), 27 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 375702264b..1989313ee6 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -66,6 +66,7 @@ __all__ = [
     "__version__",
     "HAS_FLASH_ATTENTION",
     "platform_system",
+    "patch_tokenizer",
 ]
 
 
@@ -113,6 +114,30 @@ def prepare_model_for_kbit_training(
 pass
 
 
+def patch_tokenizer(model, tokenizer):
+    if model is not None:
+        model.config.update({"unsloth_version" : __version__})
+    if not hasattr(tokenizer, "pad_token") or tokenizer.pad_token is None:
+        # Fixes https://github.com/unslothai/unsloth/issues/5
+        if hasattr(tokenizer, "unk_token"):
+            tokenizer.add_special_tokens({"pad_token" : tokenizer.unk_token})
+            tokenizer.pad_token = tokenizer.unk_token
+        else:
+            name = model.config._name_or_path if model is not None else "Model"
+            logger.warning_one(
+                f"{name} does not have a padding or unknown token!\n"\
+                f"Will use the EOS token of id {tokenizer.eos_token_id} as padding."
+            )
+            assert(hasattr(tokenizer, "eos_token"))
+            tokenizer.add_special_tokens({"pad_token" : tokenizer.eos_token})
+            tokenizer.pad_token = tokenizer.eos_token
+        if model is not None:
+            config = model.config.update({"pad_token_id" : tokenizer.eos_token_id})
+    pass
+    return model, tokenizer
+pass
+
+
 # Weirdly LoraLayer.update_layer downcasts PEFT layers to float16??
 # For mixed precision, we need it to be in float32 not float16.
 from peft.tuners.lora.layer import LoraLayer
diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py
index 3add80c6ad..1200103854 100644
--- a/unsloth/tokenizer_utils.py
+++ b/unsloth/tokenizer_utils.py
@@ -22,7 +22,6 @@ __all__ = [
     "load_correct_tokenizer",
     "fix_sentencepiece_tokenizer",
     "check_tokenizer",
-    "patch_tokenizer",
 ]
 
 
@@ -286,30 +285,6 @@ def load_correct_tokenizer(
 pass
 
 
-def patch_tokenizer(model, tokenizer):
-    if model is not None:
-        model.config.update({"unsloth_version" : __version__})
-    if not hasattr(tokenizer, "pad_token") or tokenizer.pad_token is None:
-        # Fixes https://github.com/unslothai/unsloth/issues/5
-        if hasattr(tokenizer, "unk_token"):
-            tokenizer.add_special_tokens({"pad_token" : tokenizer.unk_token})
-            tokenizer.pad_token = tokenizer.unk_token
-        else:
-            name = model.config._name_or_path if model is not None else "Model"
-            logger.warning_one(
-                f"{name} does not have a padding or unknown token!\n"\
-                f"Will use the EOS token of id {tokenizer.eos_token_id} as padding."
-            )
-            assert(hasattr(tokenizer, "eos_token"))
-            tokenizer.add_special_tokens({"pad_token" : tokenizer.eos_token})
-            tokenizer.pad_token = tokenizer.eos_token
-        if model is not None:
-            config = model.config.update({"pad_token_id" : tokenizer.eos_token_id})
-    pass
-    return model, tokenizer
-pass
-
-
 def check_tokenizer(
     model,
     tokenizer,
@@ -377,7 +352,7 @@ def check_tokenizer(
                     pass
 
                     # Recheck!
-                    can_be_removed = len(try_removal) == len(bad_tokens)
+                    can_be_removed = (len(try_removal) == len(bad_tokens))
                     if can_be_removed: remove_generic = True
                     can_be_removed1 = bad_tokens
                 pass
@@ -389,7 +364,7 @@ def check_tokenizer(
                         del tokenizer._added_tokens_decoder[remove_id]
                         del tokenizer._added_tokens_encoder[bad_token]
 
-                        if remove_generic and try_removal[j] == bad_token:
+                        if remove_generic and (try_removal[j] == bad_token):
                             # Remove sep token for example
                             setattr(tokenizer, try_mapper[j], None)
                             setattr(tokenizer, try_mapper[j] + "_id", None)