| 1423 | |
| 1424 | |
| 1425 | class VocabFactory: |
| 1426 | _VOCAB_CLASSES: list[type[Vocab]] = [SentencePieceVocab, BpeVocab, LlamaHfVocab] |
| 1427 | |
| 1428 | def __init__(self, path: Path): |
| 1429 | self.path = path |
| 1430 | |
| 1431 | def _create_special_vocab(self, vocab: BaseVocab, model_parent_path: Path) -> gguf.SpecialVocab: |
| 1432 | load_merges = vocab.name == "bpe" |
| 1433 | n_vocab = vocab.vocab_size if isinstance(vocab, Vocab) else None |
| 1434 | return gguf.SpecialVocab( |
| 1435 | model_parent_path, |
| 1436 | load_merges=load_merges, |
| 1437 | special_token_types=None, # Predetermined or passed as a parameter |
| 1438 | n_vocab=n_vocab, |
| 1439 | ) |
| 1440 | |
| 1441 | def _create_vocab_by_path(self, vocab_types: list[str]) -> Vocab: |
| 1442 | vocab_classes: dict[str, type[Vocab]] = {cls.name: cls for cls in self._VOCAB_CLASSES} |
| 1443 | selected_vocabs: dict[str, type[Vocab]] = {} |
| 1444 | for vtype in vocab_types: |
| 1445 | try: |
| 1446 | selected_vocabs[vtype] = vocab_classes[vtype] |
| 1447 | except KeyError: |
| 1448 | raise ValueError(f"Unsupported vocabulary type {vtype}") from None |
| 1449 | |
| 1450 | for vtype, cls in selected_vocabs.items(): |
| 1451 | try: |
| 1452 | vocab = cls(self.path) |
| 1453 | break |
| 1454 | except FileNotFoundError: |
| 1455 | pass # ignore unavailable tokenizers |
| 1456 | else: |
| 1457 | raise FileNotFoundError(f"Could not find a tokenizer matching any of {vocab_types}") |
| 1458 | |
| 1459 | logger.info(f"Loaded vocab file {vocab.fname_tokenizer!r}, type {vocab.name!r}") |
| 1460 | return vocab |
| 1461 | |
| 1462 | def load_vocab(self, vocab_types: list[str] | None, model_parent_path: Path) -> tuple[BaseVocab, gguf.SpecialVocab]: |
| 1463 | vocab: BaseVocab |
| 1464 | if vocab_types is None: |
| 1465 | vocab = NoVocab() |
| 1466 | else: |
| 1467 | vocab = self._create_vocab_by_path(vocab_types) |
| 1468 | # FIXME: Respect --vocab-dir? |
| 1469 | special_vocab = self._create_special_vocab( |
| 1470 | vocab, |
| 1471 | model_parent_path, |
| 1472 | ) |
| 1473 | return vocab, special_vocab |
| 1474 | |
| 1475 | |
| 1476 | def default_outfile(model_paths: list[Path], file_type: GGMLFileType) -> Path: |