278 lines
9.7 KiB
Python
278 lines
9.7 KiB
Python
# 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):
|
|
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)
|
|
if assert_same_tokenization(slow_tokenizer, fast_tokenizer):
|
|
return fast_tokenizer
|
|
return slow_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
|