| 322 | # |
| 323 | |
| 324 | class BpeVocab: |
| 325 | def __init__(self, fname_tokenizer: Path, fname_added_tokens: Path | None) -> None: |
| 326 | self.bpe_tokenizer = json.loads(open(str(fname_tokenizer), encoding="utf-8").read()) |
| 327 | added_tokens: dict[str, int] |
| 328 | if fname_added_tokens is not None: |
| 329 | # FIXME: Verify that added tokens here _cannot_ overlap with the main vocab. |
| 330 | added_tokens = json.load(open(fname_added_tokens, encoding="utf-8")) |
| 331 | else: |
| 332 | # Fall back to trying to find the added tokens in tokenizer.json |
| 333 | tokenizer_json_file = fname_tokenizer.parent / 'tokenizer.json' |
| 334 | if not tokenizer_json_file.is_file(): |
| 335 | added_tokens = {} |
| 336 | else: |
| 337 | tokenizer_json = json.load(open(tokenizer_json_file, encoding="utf-8")) |
| 338 | added_tokens = dict( |
| 339 | (item['content'], item['id']) |
| 340 | for item in tokenizer_json.get('added_tokens', []) |
| 341 | # Added tokens here can be duplicates of the main vocabulary. |
| 342 | if item['content'] not in self.bpe_tokenizer ) |
| 343 | |
| 344 | vocab_size: int = len(self.bpe_tokenizer) |
| 345 | expected_ids = list(range(vocab_size, vocab_size + len(added_tokens))) |
| 346 | actual_ids = sorted(added_tokens.values()) |
| 347 | if expected_ids != actual_ids: |
| 348 | expected_end_id = vocab_size + len(actual_ids) - 1 |
| 349 | raise Exception(f"Expected the {len(actual_ids)} added token ID(s) to be sequential in the range {vocab_size} - {expected_end_id}; got {actual_ids}") |
| 350 | |
| 351 | items = sorted(added_tokens.items(), key=lambda text_idx: text_idx[1]) |
| 352 | self.added_tokens_list = [text for (text, idx) in items] |
| 353 | self.vocab_size_base: int = vocab_size |
| 354 | self.vocab_size: int = self.vocab_size_base + len(self.added_tokens_list) |
| 355 | self.fname_tokenizer = fname_tokenizer |
| 356 | self.fname_added_tokens = fname_added_tokens |
| 357 | |
| 358 | def bpe_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]: |
| 359 | tokenizer = self.bpe_tokenizer |
| 360 | from transformers.models.gpt2 import tokenization_gpt2 |
| 361 | reverse_vocab = {id: encoded_tok for encoded_tok, id in tokenizer.items()} |
| 362 | |
| 363 | for i, _ in enumerate(tokenizer): |
| 364 | yield reverse_vocab[i], 0.0, gguf.TokenType.NORMAL |
| 365 | |
| 366 | def added_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]: |
| 367 | for text in self.added_tokens_list: |
| 368 | score = -1000.0 |
| 369 | yield text.encode("utf-8"), score, gguf.TokenType.CONTROL |
| 370 | |
| 371 | def all_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]: |
| 372 | yield from self.bpe_tokens() |
| 373 | yield from self.added_tokens() |
| 374 | |
| 375 | def __repr__(self) -> str: |
| 376 | return f"<BpeVocab with {self.vocab_size_base} base tokens and {len(self.added_tokens_list)} added tokens>" |
| 377 | |
| 378 | |
| 379 | class SentencePieceVocab: |