Heal tokenizers

unsloth/__init__.py

@@ -113,3 +113,4 @@ pass
 from .models import *
 from .save import *
 from .chat_templates import *
+from .tokenizer_utils import *

unsloth/chat_templates.py

@@ -25,6 +25,7 @@ from .models._utils import patch_tokenizer
 from .save import patch_saving_functions
 import os
 import shutil
+from .tokenizer_utils import load_correct_tokenizer, fix_sentencepiece_tokenizer
 
 CHAT_TEMPLATES = {}
 
@@ -252,66 +253,6 @@ gemma_chatml_eos_token = (
 CHAT_TEMPLATES["gemma_chatml"] = (gemma_chatml_template, gemma_chatml_eos_token,)
 
 
-def fix_sentencepiece_tokenizer(
-    old_tokenizer,
-    new_tokenizer,
-    token_mapping,
-    temporary_location = "_unsloth_sentencepiece_temp",
-):
-    # From https://github.com/google/sentencepiece/issues/121
-    # We need to manually edit the sentencepiece tokenizer!
-    try:
-        import sentencepiece.sentencepiece_model_pb2 as sentencepiece_model_pb2
-    except:
-        if not os.path.exists(temporary_location):
-            os.system("git clone https://github.com/google/sentencepiece.git unsloth_sentencepiece_temp")
-            os.system(f"cd {temporary_location}/src && protoc --python_out=. sentencepiece_model.proto")
-            shutil.rmtree(temporary_location)
-        pass
-        import sentencepiece.sentencepiece_model_pb2 as sentencepiece_model_pb2
-    pass
-
-    if not os.path.exists(temporary_location):
-        os.makedirs(temporary_location)
-    pass
-
-    # First save the old tokenizer
-    old_tokenizer.save_pretrained(temporary_location)
-
-    from sentencepiece import SentencePieceProcessor
-    tokenizer_file = sentencepiece_model_pb2.ModelProto()
-    tokenizer_file.ParseFromString(open(f"{temporary_location}/tokenizer.model", "rb").read())
-
-    # Now save the new tokenizer
-    new_tokenizer.save_pretrained(temporary_location)
-
-    # Now correct the old tokenizer's .model file
-    for old_token, new_token in token_mapping.items():
-        ids = old_tokenizer([old_token], add_special_tokens = False).input_ids
-        ids = ids[0]
-        if (len(ids) != 1):
-            # Skip this token!
-            print(f"Skip mapping {old_token} to {new_token} since {new_token} is already in the tokenizer!")
-            continue
-        pass
-        ids = ids[0]
-        tokenizer_piece = tokenizer_file.pieces[ids]
-        assert(tokenizer_piece.piece == old_token)
-        tokenizer_piece.piece = new_token
-    pass
-
-    # And now write it
-    with open(f"{temporary_location}/tokenizer.model", "wb") as file:
-        file.write(tokenizer_file.SerializeToString())
-    pass
-
-    # And load it!
-    from transformers import AutoTokenizer
-    tokenizer = AutoTokenizer.from_pretrained(temporary_location, eos_token = new_tokenizer.eos_token)
-    return tokenizer
-pass
-
-
 def get_chat_template(
     tokenizer,
     chat_template = "chatml",

unsloth/models/llama.py

@@ -26,6 +26,7 @@ from transformers.modeling_attn_mask_utils import (
 from ..kernels import *
 from ._utils import *
 from ._utils import __version__
+from ..tokenizer_utils import load_correct_tokenizer
 if HAS_FLASH_ATTENTION:
     from flash_attn import flash_attn_func
 
@@ -1014,8 +1015,8 @@ class FastLlamaModel:
 
         # Counteract saved tokenizers
         tokenizer_name = model_name if tokenizer_name is None else tokenizer_name
-        tokenizer = AutoTokenizer.from_pretrained(
-            tokenizer_name,
+        tokenizer = load_correct_tokenizer(
+            tokenizer_name    = tokenizer_name,
             model_max_length  = max_position_embeddings,
             padding_side      = "right",
             token             = token,

unsloth/models/mistral.py

@@ -362,7 +362,7 @@ class FastMistralModel(FastLlamaModel):
 
         # Counteract saved tokenizers
         tokenizer_name = model_name if tokenizer_name is None else tokenizer_name
-        tokenizer = AutoTokenizer.from_pretrained(
+        tokenizer = load_correct_tokenizer(
             tokenizer_name,
             model_max_length  = max_position_embeddings,
             padding_side      = "right",

unsloth/tokenizer_utils.py

@@ -0,0 +1,277 @@
+# 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.
+
+from transformers import AutoTokenizer
+from transformers.convert_slow_tokenizer import convert_slow_tokenizer
+from transformers import PreTrainedTokenizerFast
+import re
+import os
+
+__all__ = [
+	"load_correct_tokenizer",
+	"fix_sentencepiece_tokenizer",
+]
+
+
+def try_fix_tokenizer(tokenizer, prepend = True):
+
+    if hasattr(tokenizer, "_tokenizer"):
+        converted_tokenizer = tokenizer._tokenizer
+    else:
+        from transformers.convert_slow_tokenizer import convert_slow_tokenizer
+        converted_tokenizer = convert_slow_tokenizer(tokenizer)
+    pass
+
+    tokenizer_string = converted_tokenizer.to_str()
+
+    # Llama does ▁apple. Sometimes this is wrong!!
+    prepend_text = '{"type":"Prepend","prepend":"▁"},'
+    if not prepend and prepend_text in tokenizer_string:
+        tokenizer_string = tokenizer_string.replace(prepend_text, "", 1)
+    pass
+
+    dir_names = dir(tokenizer)
+    # Get eos_token, bos_token etc
+    token_names = [x for x in dir_names if x.endswith("_token") and x.count("_") == 1]
+
+    for token_name in token_names:
+        token = eval(f"tokenizer.{token_name}")
+        if token is None: continue
+        token_id = eval(f"tokenizer.{token_name}_id")
+
+        # Locate the token's id mapping in the string
+        find_text = f'"id":{token_id},"content":"'
+        start = tokenizer_string.find(find_text) + len(find_text)
+        if start == -1: continue
+        end   = tokenizer_string.find('",', start)
+
+        bad_token = tokenizer_string[start : end]
+        # Check if token is the actual same one - if not, edit it
+        if bad_token != token:
+            bad_text  = f'{find_text}{bad_token}",'
+            good_text = f'{find_text}{token}",'
+            tokenizer_string = tokenizer_string.replace(bad_text, good_text, 1)
+
+            # And replace vocab section
+            bad_text = f'"{bad_token}":{token_id},'
+            good_text = f'"{token}":{token_id},'
+            tokenizer_string = tokenizer_string.replace(bad_text, good_text, 1)
+        pass
+    pass
+
+    fixed_tokenizer = converted_tokenizer.from_str(tokenizer_string)
+    return fixed_tokenizer
+pass
+
+
+def get_sorted_dict(dictionary):
+	# Get sorted dict by values 0, 1, 2, ...
+    sorted_keys = sorted(dictionary.values())
+    inverted_dictionary = { value : key for key, value in dictionary.items() }
+
+    sorted_dictionary = {}
+    for key in sorted_keys:
+        value = inverted_dictionary[key]
+        sorted_dictionary[value] = key
+    return sorted_dictionary
+pass
+
+
+def convert_to_fast_tokenizer(
+    slow_tokenizer,
+    temporary_location = "_unsloth_sentencepiece_temp",
+):
+    is_fast = getattr(slow_tokenizer, "is_fast", False)
+    if is_fast: return slow_tokenizer
+    
+    try:
+        tokenizer_name = slow_tokenizer.__class__.__name__
+        lowered_tokenizer_name = tokenizer_name.lower()
+        if lowered_tokenizer_name.endswith("tokenizer"):
+            class_name = lowered_tokenizer_name[:-len("tokenizer")]
+            FastTokenizer = eval(
+                f'__import__(f"transformers.models.{class_name}").{tokenizer_name}Fast'
+            )
+        else:
+            FastTokenizer = PreTrainedTokenizerFast
+    except:
+        FastTokenizer = PreTrainedTokenizerFast
+    pass
+
+    # Get all arguments (bos_token, etc)
+    docs = FastTokenizer.__doc__
+    docs = docs[docs.find("Args:"):]
+    args = re.findall(r"\n[\s]+([^\s]{1,}) \(", docs, flags = re.MULTILINE)
+    args = [x for x in args if not x.endswith("_file")]
+
+    # Also some missing maybe!
+    docs = PreTrainedTokenizerFast.__doc__
+    docs = docs[docs.find("Args:"):]
+    args2 = re.findall(r"\n[\s]+([^\s]{1,}) \(", docs, flags = re.MULTILINE)
+    args2 = [x for x in args2 if not x.endswith("_file")]
+    args = list(set(args + args2))
+
+    kwargs = {
+        "tokenizer_object" : try_fix_tokenizer(slow_tokenizer, prepend = True),
+        "tokenizer_file"   : slow_tokenizer.vocab_file,
+    }
+    for arg in args:
+        try: kwargs[arg] = eval(f"slow_tokenizer.{arg}")
+        except: continue
+    pass
+    fast_tokenizer = FastTokenizer( **kwargs )
+
+    # Check if they're similar!
+    sorted_slow_tokenizer = get_sorted_dict(slow_tokenizer.get_vocab())
+    sorted_fast_tokenizer = get_sorted_dict(fast_tokenizer.get_vocab())
+
+    check_vocab   = (sorted_slow_tokenizer == sorted_fast_tokenizer)
+    check_special = (slow_tokenizer.all_special_tokens == fast_tokenizer.all_special_tokens)
+
+    # Failure so return slow_tokenizer
+    if not check_vocab or not check_special: return slow_tokenizer
+
+    # Now confirm if they match
+    if not assert_same_tokenization(slow_tokenizer, fast_tokenizer):
+        # Maybe remove prepending of __apple?
+        kwargs["tokenizer_object"] = try_fix_tokenizer(slow_tokenizer, prepend = False)
+        fast_tokenizer = FastTokenizer( **kwargs )
+        if not assert_same_tokenization(slow_tokenizer, fast_tokenizer):
+            # Failure :(
+            return slow_tokenizer
+        pass
+    pass
+
+    # Also tokenizer.model is missing!
+    name = slow_tokenizer.name_or_path.replace("/", "_")
+    if not os.path.exists(temporary_location):
+        os.makedirs(temporary_location)
+    pass
+    new_location = f"{temporary_location}/{name}"
+    slow_tokenizer.save_pretrained(new_location)
+    fast_tokenizer.save_pretrained(new_location)
+
+    # Now load it!
+    fast_tokenizer = AutoTokenizer.from_pretrained(new_location)
+    return fast_tokenizer
+pass
+
+
+def assert_same_tokenization(slow_tokenizer, fast_tokenizer):
+    # Get eos_token, bos_token etc
+    dir_names = dir(slow_tokenizer)
+    special_tokens = list(filter(None, (
+        eval(f"slow_tokenizer.{x}") for x in dir_names
+        if x.endswith("_token") and x.count("_") == 1
+    )))
+    all_special_tokens = list(set(special_tokens + slow_tokenizer.all_special_tokens))
+    string = "\n".join(all_special_tokens) + \
+        "A quick brown fox jumps over the lazy dog!!\n\n" + \
+        "".join(all_special_tokens)
+    return slow_tokenizer(string).input_ids == fast_tokenizer(string).input_ids
+pass
+
+
+def fix_sentencepiece_tokenizer(
+    old_tokenizer,
+    new_tokenizer,
+    token_mapping,
+    temporary_location = "_unsloth_sentencepiece_temp",
+):
+    # From https://github.com/google/sentencepiece/issues/121
+    # We need to manually edit the sentencepiece tokenizer!
+    try:
+        import sentencepiece.sentencepiece_model_pb2 as sentencepiece_model_pb2
+    except:
+        if not os.path.exists(temporary_location):
+            os.system(f"git clone https://github.com/google/sentencepiece.git {temporary_location}")
+            os.system(f"cd {temporary_location}/src && protoc --python_out=. sentencepiece_model.proto")
+            shutil.rmtree(temporary_location)
+        pass
+        import sentencepiece.sentencepiece_model_pb2 as sentencepiece_model_pb2
+    pass
+
+    if not os.path.exists(temporary_location):
+        os.makedirs(temporary_location)
+    pass
+
+    # First save the old tokenizer
+    old_tokenizer.save_pretrained(temporary_location)
+
+    from sentencepiece import SentencePieceProcessor
+    tokenizer_file = sentencepiece_model_pb2.ModelProto()
+    tokenizer_file.ParseFromString(open(f"{temporary_location}/tokenizer.model", "rb").read())
+
+    # Now save the new tokenizer
+    new_tokenizer.save_pretrained(temporary_location)
+
+    # Now correct the old tokenizer's .model file
+    for old_token, new_token in token_mapping.items():
+        ids = old_tokenizer([old_token], add_special_tokens = False).input_ids
+        ids = ids[0]
+        if (len(ids) != 1):
+            # Skip this token!
+            print(f"Skip mapping {old_token} to {new_token} since {new_token} is already in the tokenizer!")
+            continue
+        pass
+        ids = ids[0]
+        tokenizer_piece = tokenizer_file.pieces[ids]
+        assert(tokenizer_piece.piece == old_token)
+        tokenizer_piece.piece = new_token
+    pass
+
+    # And now write it
+    with open(f"{temporary_location}/tokenizer.model", "wb") as file:
+        file.write(tokenizer_file.SerializeToString())
+    pass
+
+    # And load it!
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained(temporary_location, eos_token = new_tokenizer.eos_token)
+    return tokenizer
+pass
+
+
+def load_correct_tokenizer(
+    tokenizer_name,
+    model_max_length,
+    padding_side = "right",
+    token = None,
+    trust_remote_code = False,
+):
+    slow_tokenizer = AutoTokenizer.from_pretrained(
+        tokenizer_name,
+        model_max_length  = model_max_length,
+        padding_side      = padding_side,
+        token             = token,
+        trust_remote_code = trust_remote_code,
+        use_fast          = False,
+    )
+    fast_tokenizer = AutoTokenizer.from_pretrained(
+        tokenizer_name,
+        model_max_length  = model_max_length,
+        padding_side      = padding_side,
+        token             = token,
+        trust_remote_code = trust_remote_code,
+    )
+    fast_tokenizer.add_bos_token = slow_tokenizer.add_bos_token
+    fast_tokenizer.add_eos_token = slow_tokenizer.add_eos_token
+    
+    # Confirm if slow and fast are equivalent!
+    if assert_same_tokenization(slow_tokenizer, fast_tokenizer):
+        return fast_tokenizer
+    else:
+        return slow_tokenizer
+    pass
+pass