| 377 | |
| 378 | |
| 379 | class SentencePieceVocab: |
| 380 | def __init__(self, fname_tokenizer: Path, fname_added_tokens: Path | None) -> None: |
| 381 | self.sentencepiece_tokenizer = SentencePieceProcessor(str(fname_tokenizer)) |
| 382 | added_tokens: dict[str, int] |
| 383 | if fname_added_tokens is not None: |
| 384 | added_tokens = json.load(open(fname_added_tokens, encoding="utf-8")) |
| 385 | else: |
| 386 | added_tokens = {} |
| 387 | |
| 388 | vocab_size: int = self.sentencepiece_tokenizer.vocab_size() |
| 389 | |
| 390 | new_tokens = {id: piece for piece, id in added_tokens.items() if id >= vocab_size} |
| 391 | expected_new_ids = list(range(vocab_size, vocab_size + len(new_tokens))) |
| 392 | actual_new_ids = sorted(new_tokens.keys()) |
| 393 | |
| 394 | if expected_new_ids != actual_new_ids: |
| 395 | raise ValueError(f"Expected new token IDs {expected_new_ids} to be sequential; got {actual_new_ids}") |
| 396 | |
| 397 | # Token pieces that were added to the base vocabulary. |
| 398 | self.added_tokens_list = [new_tokens[id] for id in actual_new_ids] |
| 399 | self.vocab_size_base = vocab_size |
| 400 | self.vocab_size = self.vocab_size_base + len(self.added_tokens_list) |
| 401 | self.fname_tokenizer = fname_tokenizer |
| 402 | self.fname_added_tokens = fname_added_tokens |
| 403 | |
| 404 | def sentencepiece_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]: |
| 405 | tokenizer = self.sentencepiece_tokenizer |
| 406 | for i in range(tokenizer.vocab_size()): |
| 407 | piece = tokenizer.id_to_piece(i) |
| 408 | text: bytes = piece.encode("utf-8") |
| 409 | score: float = tokenizer.get_score(i) |
| 410 | |
| 411 | toktype = gguf.TokenType.NORMAL |
| 412 | if tokenizer.is_unknown(i): |
| 413 | toktype = gguf.TokenType.UNKNOWN |
| 414 | if tokenizer.is_control(i): |
| 415 | toktype = gguf.TokenType.CONTROL |
| 416 | |
| 417 | # NOTE: I think added_tokens are user defined. |
| 418 | # ref: https://github.com/google/sentencepiece/blob/master/src/sentencepiece_model.proto |
| 419 | # if tokenizer.is_user_defined(i): toktype = gguf.TokenType.USER_DEFINED |
| 420 | |
| 421 | if tokenizer.is_unused(i): |
| 422 | toktype = gguf.TokenType.UNUSED |
| 423 | if tokenizer.is_byte(i): |
| 424 | toktype = gguf.TokenType.BYTE |
| 425 | |
| 426 | yield text, score, toktype |
| 427 | |
| 428 | def added_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]: |
| 429 | for text in self.added_tokens_list: |
| 430 | score = -1000.0 |
| 431 | yield text.encode("utf-8"), score, gguf.TokenType.USER_DEFINED |
| 432 | |
| 433 | def all_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]: |
| 434 | yield from self.sentencepiece_tokens() |
| 435 | yield from self.added_tokens() |
| 436 | |