| 1285 | |
| 1286 | |
| 1287 | class VocabFactory: |
| 1288 | _FILES = {"spm": "tokenizer.model", "bpe": "vocab.json", "hfft": "tokenizer.json"} |
| 1289 | |
| 1290 | def __init__(self, path: Path): |
| 1291 | self.path = path |
| 1292 | self.file_paths = self._detect_files() |
| 1293 | print(f"Found vocab files: {self.file_paths}") |
| 1294 | |
| 1295 | def _detect_files(self) -> dict[str, Path | None]: |
| 1296 | def locate(file: str) -> Path | None: |
| 1297 | if (path := self.path / file).exists(): |
| 1298 | return path |
| 1299 | if (path := self.path.parent / file).exists(): |
| 1300 | return path |
| 1301 | return None |
| 1302 | |
| 1303 | return {vt: locate(f) for vt, f in self._FILES.items()} |
| 1304 | |
| 1305 | def _select_file(self, vocab_types: list[str]) -> tuple[str, Path]: |
| 1306 | for vtype in vocab_types: |
| 1307 | try: |
| 1308 | path = self.file_paths[vtype] |
| 1309 | except KeyError: |
| 1310 | raise ValueError(f"Unsupported vocabulary type {vtype}") from None |
| 1311 | if path is not None: |
| 1312 | return vtype, path |
| 1313 | raise FileNotFoundError(f"Could not find any of {[self._FILES[vt] for vt in vocab_types]}") |
| 1314 | |
| 1315 | def _create_special_vocab(self, vocab: Vocab, vocabtype: str, model_parent_path: Path) -> gguf.SpecialVocab: |
| 1316 | load_merges = vocabtype == "bpe" |
| 1317 | n_vocab = vocab.vocab_size if hasattr(vocab, "vocab_size") else None |
| 1318 | return gguf.SpecialVocab( |
| 1319 | model_parent_path, |
| 1320 | load_merges=load_merges, |
| 1321 | special_token_types=None, # Predetermined or passed as a parameter |
| 1322 | n_vocab=n_vocab, |
| 1323 | ) |
| 1324 | |
| 1325 | def load_vocab(self, vocab_types: list[str], model_parent_path: Path) -> tuple[Vocab, gguf.SpecialVocab]: |
| 1326 | vocab_type, path = self._select_file(vocab_types) |
| 1327 | print(f"Loading vocab file {path!r}, type {vocab_type!r}") |
| 1328 | |
| 1329 | added_tokens_path = path.parent / "added_tokens.json" |
| 1330 | vocab: Vocab |
| 1331 | if vocab_type == "bpe": |
| 1332 | vocab = BpeVocab( |
| 1333 | path, added_tokens_path if added_tokens_path.exists() else None |
| 1334 | ) |
| 1335 | elif vocab_type == "spm": |
| 1336 | vocab = SentencePieceVocab( |
| 1337 | path, added_tokens_path if added_tokens_path.exists() else None |
| 1338 | ) |
| 1339 | elif vocab_type == "hfft": |
| 1340 | vocab = HfVocab( |
| 1341 | path.parent, added_tokens_path if added_tokens_path.exists() else None |
| 1342 | ) |
| 1343 | else: |
| 1344 | raise ValueError(vocab_type) |