From a330a53623a28152e3441dfaebd26dbfcb09886a Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Sun, 24 Mar 2024 04:20:30 +1100 Subject: [PATCH] Heal tokenizers --- unsloth/__init__.py | 1 + unsloth/chat_templates.py | 61 +------- unsloth/models/llama.py | 5 +- unsloth/models/mistral.py | 2 +- unsloth/tokenizer_utils.py | 277 +++++++++++++++++++++++++++++++++++++ 5 files changed, 283 insertions(+), 63 deletions(-) create mode 100644 unsloth/tokenizer_utils.py diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 7d271edf4e..d4ca45d7d1 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -113,3 +113,4 @@ pass from .models import * from .save import * from .chat_templates import * +from .tokenizer_utils import * diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index 939ed7210d..e8577fa6c8 100644 --- a/unsloth/chat_templates.py +++ b/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", diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 6b258dbb6b..d3b48d9439 100644 --- a/unsloth/models/llama.py +++ b/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, diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index f650d8f12e..9c73266575 100644 --- a/unsloth/models/mistral.py +++ b/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", diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py new file mode 100644 index 0000000000..60c8ceb6fd --- /dev/null +++ b/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