test-tokenizer-random.py
1 # Test libllama tokenizer == AutoTokenizer. 2 # Brute force random words/text generation. 3 # 4 # Sample usage: 5 # 6 # python3 tests/test-tokenizer-random.py ./models/ggml-vocab-llama-bpe.gguf ./models/tokenizers/llama-bpe 7 # 8 9 import time 10 import logging 11 import argparse 12 import subprocess 13 import random 14 15 from typing import Callable, Iterator 16 17 import cffi 18 from transformers import AutoTokenizer 19 20 logger = logging.getLogger("test-tokenizer-random-bpe") 21 22 23 class LibLlama: 24 25 DEFAULT_PATH_LLAMA_H = "./llama.h" 26 DEFAULT_PATH_LIBLLAMA = "./build/libllama.so" # CMakeLists.txt: BUILD_SHARED_LIBS ON 27 28 def __init__(self, path_llama_h: str = None, path_libllama: str = None): 29 path_llama_h = path_llama_h or self.DEFAULT_PATH_LLAMA_H 30 path_libllama = path_libllama or self.DEFAULT_PATH_LIBLLAMA 31 (self.ffi, self.lib) = self._load_libllama_cffi(path_llama_h, path_libllama) 32 self.lib.llama_backend_init() 33 34 def _load_libllama_cffi(self, path_llama_h: str, path_libllama: str): 35 cmd = ["gcc", "-E", "-P", "-D__restrict=", "-D__attribute__(x)=", "-D__asm__(x)=", path_llama_h] 36 res = subprocess.run(cmd, stdout=subprocess.PIPE) 37 assert (res.returncode == 0) 38 source = res.stdout.decode() 39 ffi = cffi.FFI() 40 if True: # workarounds for pycparser 41 source = "typedef struct { } __builtin_va_list;" + "\n" + source 42 source = source.replace("sizeof (int)", str(ffi.sizeof("int"))) 43 source = source.replace("sizeof (void *)", str(ffi.sizeof("void*"))) 44 source = source.replace("sizeof (size_t)", str(ffi.sizeof("size_t"))) 45 source = source.replace("sizeof(int32_t)", str(ffi.sizeof("int32_t"))) 46 ffi.cdef(source, override=True) 47 lib = ffi.dlopen(path_libllama) 48 return (ffi, lib) 49 50 def model_default_params(self, **kwargs): 51 mparams = self.lib.llama_model_default_params() 52 for k, v in kwargs.items(): 53 setattr(mparams, k, v) 54 return mparams 55 56 def context_default_params(self, **kwargs): 57 cparams = self.lib.llama_context_default_params() 58 for k, v in kwargs.items(): 59 setattr(cparams, k, v) 60 return cparams 61 62 63 class LibLlamaModel: 64 65 def __init__(self, libllama: LibLlama, path_model: str, mparams={}, cparams={}): 66 self.lib = libllama.lib 67 self.ffi = libllama.ffi 68 if isinstance(mparams, dict): 69 mparams = libllama.model_default_params(**mparams) 70 self.model = self.lib.llama_load_model_from_file(path_model.encode(), mparams) 71 if not self.model: 72 raise RuntimeError("error: failed to load model '%s'" % path_model) 73 if isinstance(cparams, dict): 74 cparams = libllama.context_default_params(**cparams) 75 self.ctx = self.lib.llama_new_context_with_model(self.model, cparams) 76 if not self.ctx: 77 raise RuntimeError("error: failed to create context for model '%s'" % path_model) 78 n_tokens_max = self.lib.llama_n_ctx(self.ctx) 79 self.token_ids = self.ffi.new("llama_token[]", n_tokens_max) 80 81 def free(self): 82 if self.ctx: 83 self.lib.llama_free(self.ctx) 84 if self.model: 85 self.lib.llama_free_model(self.model) 86 self.ctx = None 87 self.model = None 88 self.lib = None 89 90 def tokenize(self, text: str, n_tokens_max: int = 0, add_special: bool = False, parse_special: bool = False) -> list[int]: 91 n_tokens_max = n_tokens_max if n_tokens_max > 0 else len(self.token_ids) 92 text = text.encode("utf-8") 93 num = self.lib.llama_tokenize(self.model, text, len(text), self.token_ids, n_tokens_max, add_special, parse_special) 94 if num < 0: 95 return [] 96 return list(self.token_ids[0:num]) 97 98 99 def generator_custom_text() -> Iterator[str]: 100 """General tests""" 101 yield from [ 102 "", 103 " ", 104 " ", 105 " ", 106 "\t", 107 "\n", 108 "\n\n", 109 "\n\n\n", 110 "\t\n", 111 "Hello world", 112 " Hello world", 113 "Hello World", 114 " Hello World", 115 " Hello World!", 116 "Hello, world!", 117 " Hello, world!", 118 " this is 🦙.cpp", 119 "w048 7tuijk dsdfhu", 120 "нещо на Български", 121 "កាន់តែពិសេសអាចខលចេញ", 122 "🚀 (normal) 😶🌫️ (multiple emojis concatenated) ✅ (only emoji that has its own token)", 123 "Hello", 124 " Hello", 125 " Hello", 126 " Hello", 127 " Hello", 128 " Hello\n Hello", 129 " (", 130 "\n =", 131 "' era", 132 "Hello, y'all! How are you 😁 ?我想在apple工作1314151天～", 133 "3", 134 "33", 135 "333", 136 "3333", 137 "33333", 138 "333333", 139 "3333333", 140 "33333333", 141 "333333333", 142 ] 143 144 145 def generator_custom_text_edge_cases() -> Iterator[str]: 146 """Edge cases found while debugging""" 147 yield from [ 148 '\x1f-a', # unicode_ranges_control, {0x00001C, 0x00001F} 149 '¼-a', # unicode_ranges_digit, 0x00BC 150 '½-a', # unicode_ranges_digit, 0x00BD 151 '¾-a', # unicode_ranges_digit, 0x00BE 152 'a 〇b', # unicode_ranges_digit, 0x3007 153 'Ⅵ-a', # unicode_ranges_digit, {0x00002150, 0x0000218F} // Number Forms 154 '\uFEFF//', # unicode_ranges_control, 0xFEFF (BOM) 155 'Cửa Việt', # llama-3, ignore_merges = true 156 '<s>a', # Phi-3 fail 157 '<unk><|endoftext|><s>', # Phi-3 fail 158 'a\na', # TODO: Bert fail 159 'a </s> b', # rstrip phi-3 160 'a <mask> b', # lstrip jina-v2 161 ] 162 163 164 def generator_vocab_words(vocab: list[str]) -> Iterator[str]: 165 """Brute force check all vocab words""" 166 yield from vocab 167 168 169 def generator_added_lr_strip(tokenizer) -> Iterator[str]: 170 WHITESPACES = ["", " ", " ", " "] 171 special_tokens = list(tokenizer.all_special_tokens) 172 added_tokens = list(tokenizer.added_tokens_encoder) 173 all_tokens = list(sorted(set(special_tokens + added_tokens))) 174 for token in all_tokens: 175 for lstrip in WHITESPACES: 176 for rstrip in WHITESPACES: 177 yield lstrip + token + rstrip 178 yield "a" + lstrip + token + rstrip 179 yield lstrip + token + rstrip + "z" 180 yield "a" + lstrip + token + rstrip + "z" 181 182 183 def generator_random_added_tokens(tokenizer, iterations=100) -> Iterator[str]: 184 special_tokens = list(tokenizer.all_special_tokens) 185 added_tokens = list(tokenizer.added_tokens_encoder) 186 separations = [" ", "\n", "\t", "-", "!", "one", "1", "<s>", "</s>"] 187 all_tokens = list(sorted(set(special_tokens + added_tokens + separations))) 188 rand = random.Random() 189 for m in range(iterations): 190 rand.seed(m) 191 words = rand.choices(all_tokens, k=500) 192 if words[0] == tokenizer.bos_token: # skip spam warning of double BOS 193 while len(words) > 1 and words[1] == tokenizer.bos_token: # leave one starting BOS 194 words.pop(0) 195 if tokenizer.add_bos_token: # drop all starting BOS 196 words.pop(0) 197 yield "".join(words) 198 199 200 def generator_random_chars(iterations=100) -> Iterator[str]: 201 """Brute force random text with simple characters""" 202 203 WHITESPACES = list(" " * 20 + "\n" * 5 + "\r\n" * 5 + "\t" * 5) 204 CHARS = list(sorted(set(""" 205 ABCDEFGHIJKLMNOPQRSTUVWXYZ 206 abcdefghijklmnopqrstuvwxyz 207 ÁÉÍÓÚÀÈÌÒÙÂÊÎÔÛÄËÏÖÜ 208 áéíóúàèìòùâêîôûäëïöü 209 .-,*/-+ª!"·$%&/()=?¿[]{}<>\\|@#~½¬~;:_ 210 """))) 211 212 rand = random.Random() 213 for m in range(iterations): 214 rand.seed(m) 215 text = [] 216 num_words = rand.randint(300, 400) 217 for i in range(num_words): 218 k = rand.randint(1, 7) 219 word = rand.choices(CHARS, k=k) 220 space = rand.choice(WHITESPACES) 221 text.append("".join(word) + space) 222 yield "".join(text) 223 224 225 def generator_random_vocab_chars(vocab: list[str], iterations=100) -> Iterator[str]: 226 """Brute force random text with vocab characters""" 227 228 vocab_chars = set() 229 for word in vocab: 230 vocab_chars.update(word) 231 vocab_chars = list(sorted(vocab_chars)) 232 233 rand = random.Random() 234 for m in range(iterations): 235 rand.seed(m) 236 text = rand.choices(vocab_chars, k=1024) 237 yield "".join(text) 238 239 240 def generator_random_vocab_words(vocab: list[str], iterations=100) -> Iterator[str]: 241 """Brute force random text from vocab words""" 242 243 vocab = [w.strip() for w in vocab] 244 yield from vocab 245 246 rand = random.Random() 247 for m in range(iterations): 248 rand.seed(m) 249 text = [] 250 num_words = rand.randint(300, 400) 251 for i in range(num_words): 252 k = rand.randint(1, 3) 253 words = rand.choices(vocab, k=k) 254 sep = rand.choice(" \n\r\t") 255 text.append("".join(words) + sep) 256 yield "".join(text) 257 258 259 def generator_random_bytes(iterations=100) -> Iterator[str]: 260 """Brute force random bytes""" 261 262 WHITESPACES = list(" " * 20 + "\n" * 5 + "\r\n" * 5 + "\t" * 5) 263 264 rand = random.Random() 265 for m in range(iterations): 266 rand.seed(m) 267 text = [] 268 num_words = rand.randint(300, 400) 269 for i in range(num_words): 270 k = rand.randint(1, 8) 271 word = [chr(r) for r in rand.randbytes(k) if r] 272 word.append(rand.choice(WHITESPACES)) 273 text.append("".join(word)) 274 yield "".join(text) 275 276 277 def test_compare_tokenizer(func_tokenize1: Callable, func_tokenize2: Callable, generator: Iterator[str]): 278 279 def find_first_mismatch(ids1: list[int], ids2: list[int]): 280 for i, (a, b) in enumerate(zip(ids1, ids2)): 281 if a != b: 282 return i 283 if len(ids1) == len(ids2): 284 return -1 285 return min(len(ids1), len(ids2)) 286 287 t0 = time.perf_counter() 288 logger.info("%s: %s" % (generator.__name__, "ini")) 289 for text in generator: 290 ids1 = func_tokenize1(text) 291 ids2 = func_tokenize2(text) 292 if ids1 != ids2: 293 i = find_first_mismatch(ids1, ids2) 294 ids1 = list(ids1)[max(0, i - 2) : i + 5 + 1] 295 ids2 = list(ids2)[max(0, i - 2) : i + 5 + 1] 296 logger.info(" TokenIDs: " + str(ids1)) 297 logger.info(" Expected: " + str(ids2)) 298 raise Exception() 299 t1 = time.perf_counter() 300 logger.info("%s: end, time: %.3f secs" % (generator.__name__, t1 - t0)) 301 302 303 def main(argv: list[str] = None): 304 parser = argparse.ArgumentParser() 305 parser.add_argument("vocab_file", help="path to vocab 'gguf' file") 306 parser.add_argument("dir_tokenizer", help="directory containing 'tokenizer.model' file") 307 parser.add_argument("--verbose", action="store_true", help="increase output verbosity") 308 args = parser.parse_args(argv) 309 310 logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO) 311 312 model = LibLlamaModel(LibLlama(), args.vocab_file, mparams=dict(vocab_only=True), cparams=dict(n_ctx=4096)) 313 tokenizer = AutoTokenizer.from_pretrained(args.dir_tokenizer) 314 315 def func_tokenize1(text: str): 316 return model.tokenize(text, add_special=True, parse_special=True) 317 318 def func_tokenize2(text: str): 319 return tokenizer.encode(text, add_special_tokens=True) 320 321 ids = func_tokenize2("a") 322 assert 1 <= len(ids) <= 3 323 add_bos_token = len(ids) > 1 and tokenizer.bos_token_id == ids[0] 324 tokenizer.add_bos_token = getattr(tokenizer, "add_bos_token", add_bos_token) 325 326 vocab = list(sorted(tokenizer.batch_decode(list(tokenizer.get_vocab().values()), skip_special_tokens=True))) 327 test_compare_tokenizer(func_tokenize1, func_tokenize2, generator_custom_text()) 328 test_compare_tokenizer(func_tokenize1, func_tokenize2, generator_custom_text_edge_cases()) 329 test_compare_tokenizer(func_tokenize1, func_tokenize2, generator_vocab_words(vocab)) 330 test_compare_tokenizer(func_tokenize1, func_tokenize2, generator_added_lr_strip(tokenizer)) 331 test_compare_tokenizer(func_tokenize1, func_tokenize2, generator_random_added_tokens(tokenizer, 10_000)) 332 test_compare_tokenizer(func_tokenize1, func_tokenize2, generator_random_chars(10_000)) 333 test_compare_tokenizer(func_tokenize1, func_tokenize2, generator_random_vocab_chars(vocab, 10_000)) 334 test_compare_tokenizer(func_tokenize1, func_tokenize2, generator_random_vocab_words(vocab, 5_000)) 335 # test_compare_tokenizer(func_tokenize1, func_tokenize2, generator_random_bytes(10_000)) # FAIL 336 337 model.free() 338 339 340 if __name__ == "__main__": 341 # main() 342 343 path_tokenizers = "./models/tokenizers/" 344 path_vocab_format = "./models/ggml-vocab-%s.gguf" 345 346 # import os 347 # tokenizers = os.listdir(path_tokenizers) 348 tokenizers = [ 349 "llama-spm", # SPM 350 "phi-3", # SPM 351 "jina-v2-en", # WPM 352 "bert-bge", # WPM 353 ] 354 355 for tokenizer in tokenizers: 356 print("\n" + "=" * 50 + "\n" + tokenizer + "\n") # noqa 357 vocab_file = path_vocab_format % tokenizer 358 dir_tokenizer = path_tokenizers + "/" + tokenizer 359 main([vocab_file, dir_tokenizer, "--verbose"])