hey! this is a static git repo viewer. commit diffs, file history, older commits and other revisions won't load without javascript.

you are still able to view source files on branch/tag tips, but the experience is degraded.

char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

o7sivocab : adopt leading TemplateProcessing special token as BOS (#24428)1bfbdb1342026-06-11

1from __future__ import annotations
2
3from enum import Enum
4import re
5import logging
6import json
7import os
8from pathlib import Path
9from typing import Any, Callable, Sequence, Mapping, Iterable, Protocol, ClassVar, runtime_checkable
10
11try:
12    from sentencepiece import SentencePieceProcessor
13except ImportError:
14    SentencePieceProcessor: Any = None
15
16try:
17    from mistral_common.tokens.tokenizers.mistral import MistralTokenizer # type: ignore[import-not-found, ty:unresolved-import]
18    from mistral_common.tokens.tokenizers.tekken import Tekkenizer # type: ignore[import-not-found, ty:unresolved-import]
19    from mistral_common.tokens.tokenizers.utils import ( # type: ignore[import-not-found, ty:unresolved-import]
20        _filter_valid_tokenizer_files,
21    )
22    from mistral_common.tokens.tokenizers.sentencepiece import ( # type: ignore[import-not-found, ty:unresolved-import]
23        SentencePieceTokenizer,
24    )
25except ImportError:
26    _mistral_common_installed = False
27    MistralTokenizer: Any = None
28    Tekkenizer: Any = None
29    SentencePieceTokenizer: Any = None
30    _filter_valid_tokenizer_files: Any = None
31else:
32    _mistral_common_installed = True
33
34try:
35    from mistral_common.tokens.tokenizers.utils import ( # type: ignore[import-not-found, ty:unresolved-import]
36        get_one_valid_tokenizer_file,
37    )
38except ImportError:
39    # We still want the conversion to work with older mistral-common versions.
40    get_one_valid_tokenizer_file: Any = None
41
42
43import gguf
44
45from .gguf_writer import GGUFWriter
46
47logger = logging.getLogger(__name__)
48
49
50class SpecialVocab:
51    merges: list[str]
52    add_special_token: dict[str, bool]
53    special_token_ids: dict[str, int]
54    chat_template: str | Sequence[Mapping[str, str]] | None
55    normalizer_lowercase: bool | None
56    normalizer_strip_accents: bool | None
57
58    def __init__(
59        self, path: str | os.PathLike[str], load_merges: bool = False,
60        special_token_types: Iterable[str] | None = None,
61        n_vocab: int | None = None,
62    ):
63        self.special_token_ids = {}
64        self.add_special_token = {}
65        self.n_vocab = n_vocab
66        self.load_merges = load_merges
67        self.merges = []
68        self.chat_template = None
69        self.normalizer_lowercase = None
70        self.normalizer_strip_accents = None
71        if special_token_types is not None:
72            self.special_token_types = special_token_types
73        else:
74            self.special_token_types = ('bos', 'eos', 'unk', 'sep', 'pad', 'cls', 'mask')
75        self._load(Path(path))
76
77    def __repr__(self) -> str:
78        return '<SpecialVocab with {} merges, special tokens {}, add special tokens {}>'.format(
79            len(self.merges), self.special_token_ids or "unset", self.add_special_token or "unset",
80        )
81
82    def add_to_gguf(self, gw: GGUFWriter, quiet: bool = False) -> None:
83        if self.merges:
84            if not quiet:
85                logger.info(f'Adding {len(self.merges)} merge(s).')
86            gw.add_token_merges(self.merges)
87        elif self.load_merges:
88            logger.warning('Adding merges requested but no merges found, output may be non-functional.')
89        for typ, tokid in self.special_token_ids.items():
90            id_handler: Callable[[int], None] | None = getattr(gw, f'add_{typ}_token_id', None)
91            if id_handler is None:
92                logger.warning(f'No handler for special token type {typ} with id {tokid} - skipping')
93                continue
94            if not quiet:
95                logger.info(f'Setting special token type {typ} to {tokid}')
96            id_handler(tokid)
97        for typ, value in self.add_special_token.items():
98            add_handler: Callable[[bool], None] | None = getattr(gw, f'add_add_{typ}_token', None)
99            if add_handler is None:
100                logger.warning(f'No handler for add_{typ}_token with value {value} - skipping')
101                continue
102            if not quiet:
103                logger.info(f'Setting add_{typ}_token to {value}')
104            add_handler(value)
105        if self.chat_template is not None:
106            if not quiet:
107                logger.info(f'Setting chat_template to {self.chat_template}')
108            gw.add_chat_template(self.chat_template)
109        if self.normalizer_lowercase is not None:
110            if not quiet:
111                logger.info(f'Setting normalizer_lowercase to {self.normalizer_lowercase}')
112            gw.add_normalizer_lowercase(self.normalizer_lowercase)
113        if self.normalizer_strip_accents is not None:
114            if not quiet:
115                logger.info(f'Setting normalizer_strip_accents to {self.normalizer_strip_accents}')
116            gw.add_normalizer_strip_accents(self.normalizer_strip_accents)
117
118    def _load(self, path: Path) -> None:
119        self._try_load_from_tokenizer_json(path)
120        self._try_load_from_config_json(path)
121        if self.load_merges and not self.merges:
122            self._try_load_merges_txt(path)
123
124    def _try_load_merges_txt(self, path: Path) -> bool:
125        merges_file = path / 'merges.txt'
126        if not merges_file.is_file():
127            return False
128        with open(merges_file, 'r', encoding = 'utf-8') as fp:
129            first_line = next(fp, '').strip()
130            if not first_line.startswith('#'):
131                fp.seek(0)
132                line_num = 0
133            else:
134                line_num = 1
135            merges = []
136            for line in fp:
137                line_num += 1
138                line = line.strip()
139                if not line:
140                    continue
141                parts = line.split(None, 3)
142                if len(parts) != 2:
143                    logger.warning(f'{merges_file.name}: Line {line_num}: Entry malformed, ignoring')
144                    continue
145                merges.append(f'{parts[0]} {parts[1]}')
146        self.merges = merges
147        return True
148
149    def _set_special_token(self, typ: str, tid: Any) -> None:
150        if not isinstance(tid, int):
151            return
152        if tid < 0:
153            raise ValueError(f'invalid value for special token type {typ}: {tid}')
154        if self.n_vocab is None or tid < self.n_vocab:
155            if typ in self.special_token_ids:
156                return
157            self.special_token_ids[typ] = tid
158            return
159        logger.warning(f'Special token type {typ}, id {tid} out of range, must be under {self.n_vocab} - skipping')
160
161    def _parse_normalizer(self, normalizer: dict) -> None:
162        # ref: https://huggingface.co/docs/tokenizers/api/normalizers
163        #
164        # Extracts normalizer flags from three possible formats:
165        # 1. Standalone:           {"type": "Lowercase"}
166        # 2. BertNormalizer attrs: {"type": "BertNormalizer", ...}
167        # 3. Nested in Sequence:   {"type": "Sequence", "normalizers": [...]}
168
169        normalizer_type = normalizer.get('type')
170        if normalizer_type == 'Lowercase':
171            self.normalizer_lowercase = True
172        elif normalizer_type == 'StripAccents':
173            self.normalizer_strip_accents = True
174        elif normalizer_type == 'BertNormalizer':
175            if 'lowercase' in normalizer:
176                self.normalizer_lowercase = normalizer['lowercase']
177            if 'strip_accents' in normalizer:
178                self.normalizer_strip_accents = normalizer['strip_accents']
179        elif normalizer_type == 'Sequence':
180            for norm in normalizer.get('normalizers', []):
181                self._parse_normalizer(norm)
182
183    def _try_load_from_tokenizer_json(self, path: Path) -> bool:
184        tokenizer = None
185        tokenizer_file = path / 'tokenizer.json'
186        if tokenizer_file.is_file():
187            with open(tokenizer_file, encoding = 'utf-8') as f:
188                tokenizer = json.load(f)
189            if self.load_merges:
190                merges = tokenizer.get('model', {}).get('merges')
191                if isinstance(merges, list) and merges:
192                    if isinstance(merges[0], str):
193                        self.merges = merges
194                    elif isinstance(merges[0], list) and len(merges[0]) == 2 and isinstance(merges[0][0], str):
195                        # New format since transformers 4.45 to support spaces in merges
196                        # ref: https://github.com/ggml-org/llama.cpp/issues/9692
197                        # TODO: internally store as the new format instead of converting to old
198                        if any(' ' in s for pair in merges for s in pair):
199                            logger.warning(f'Spaces in merges detected, encoding as {chr(ord(" ") + 256)!r}')
200                        self.merges = [
201                            ' '.join(
202                                [
203                                    # ensure the spaces are properly encoded
204                                    ''.join(
205                                        chr(ord(c) + 256) if c == ' ' else c
206                                        for c in part
207                                    )
208                                    for part in pair
209                                ]
210                            )
211                            for pair in merges
212                        ]
213                    else:
214                        raise ValueError("Unknown tokenizer merges format")
215            # Parse normalizer configuration (e.g. Lowercase) into metadata
216            if normalizer := tokenizer.get('normalizer'):
217                self._parse_normalizer(normalizer)
218            added_tokens = tokenizer.get('added_tokens', {})
219        else:
220            added_tokens = {}
221        tokenizer_config = None
222        tokenizer_config_file = path / 'tokenizer_config.json'
223        if tokenizer_config_file.is_file():
224            with open(tokenizer_config_file, encoding = 'utf-8') as f:
225                tokenizer_config = json.load(f)
226        if tokenizer:
227            special_bos = (tokenizer_config or {}).get('bos_token')
228            special_cls = (tokenizer_config or {}).get('cls_token')
229            special_eos = (tokenizer_config or {}).get('eos_token')
230            special_sep = (tokenizer_config or {}).get('sep_token')
231            if not special_bos and special_cls and tokenizer_config:
232                tokenizer_config['bos_token'] = special_bos = special_cls
233            if not special_eos and special_sep and tokenizer_config:
234                tokenizer_config['eos_token'] = special_eos = special_sep
235            if post_processor := tokenizer.get('post_processor'):
236                for processor in post_processor.get('processors', [post_processor]):
237                    if processor.get('type') == 'RobertaProcessing':
238                        self.add_special_token['bos'] = True
239                        self.add_special_token['eos'] = True
240                        self.add_special_token['sep'] = True
241                        if not special_cls and tokenizer_config:
242                            special_cls = processor.get('cls', [special_bos])[0]
243                            tokenizer_config['cls_token'] = special_cls
244                        if not special_sep and tokenizer_config:
245                            special_sep = processor.get('sep', [special_eos])[0]
246                            tokenizer_config['sep_token'] = special_sep
247                        continue
248                    # Crude parsing of TemplateProcessing to determine if BOS/SEP/EOS should be added
249                    # Only works with simple templates, **will** get it wrong on unusual sequences
250                    if processor.get('type') == 'TemplateProcessing':
251                        tmpl_single = processor.get('single', [])
252                        tmpl_pair = processor.get('pair', [])
253                        special_first = None
254                        special_last = None
255                        if len(tmpl_single) > 1:
256                            if special_first := tmpl_single[0].get('SpecialToken', {}).get('id'):
257                                if not tokenizer_config:
258                                    special_bos = special_first
259                                elif special_first not in (special_bos, special_cls):
260                                    if not special_bos:
261                                        tokenizer_config['bos_token'] = special_bos = special_first
262                                    if not special_cls:
263                                        tokenizer_config['cls_token'] = special_cls = special_first
264                                self.add_special_token['bos'] = True if special_first in (special_bos, special_cls) else False
265                                if special_first not in (special_bos, special_cls):
266                                    logger.warning(f'Unknown leading special token {special_first!r} in TemplateProcessing<single>')
267                            if special_last := tmpl_single[-1].get('SpecialToken', {}).get('id'):
268                                if not tokenizer_config:
269                                    special_eos = special_last
270                                elif special_last != special_eos:
271                                    if 'eot' not in self.special_token_types:
272                                        self.special_token_types = tuple(self.special_token_types) + ('eot', )
273                                        tokenizer_config['eot_token'] = special_eos
274                                    elif 'eom' not in self.special_token_types:
275                                        self.special_token_types = tuple(self.special_token_types) + ('eom', )
276                                        tokenizer_config['eom_token'] = special_eos
277                                    else:
278                                        logger.warning(f'Overriding EOS token {special_eos!r} with {special_last!r} without EOT/EOM fallback!')
279                                    tokenizer_config['eos_token'] = special_eos = special_last
280                                self.add_special_token['eos'] = True if special_last == special_eos else False
281                                if special_last != special_eos:
282                                    logger.warning(f'Unknown trailing special token {special_last!r} in TemplateProcessing<single>')
283                        if tmpl_pair:
284                            seq_start = 1 if special_first and tmpl_pair[0].get('SpecialToken', {}).get('id') == special_first else 0
285                            seq_stop = -1 if special_last and tmpl_pair[-1].get('SpecialToken', {}).get('id') == special_last else None
286                            if (special_first and seq_start == 0) or (special_last and seq_stop is None):
287                                logger.warning('TemplateProcessing<single> leading/trailing special tokens do not match TemplateProcessing<pair>')
288                            if tmpl_pair := tmpl_pair[slice(seq_start, seq_stop)]:
289                                tmpl_a = tmpl_pair[0].get('Sequence', {}).get('id')
290                                tmpl_b = tmpl_pair[-1].get('Sequence', {}).get('id')
291                                if tmpl_a != 'A' or tmpl_b != 'B':
292                                    logger.warning(f'Unknown sequence {tmpl_a}...{tmpl_b} in TemplateProcessing<pair>')
293                                # A [sep] [eos] B
294                                if tmpl_a == 'A' and tmpl_b == 'B' and (tmpl_pair := tmpl_pair[1:-1]):
295                                    add_sep = False
296                                    if special_entry := tmpl_pair[0].get('SpecialToken', {}).get('id'):
297                                        if special_entry in (special_sep, special_eos) and not special_last:
298                                            add_sep = True
299                                        if special_entry not in (special_sep, special_eos):
300                                            logger.warning(f'Unknown separator token {special_entry!r} in TemplateProcessing<pair>')
301                                    else:
302                                        logger.warning(f'Unknown middle sequence {tmpl_pair[0]!r} in TemplateProcessing<pair>')
303                                    if len(tmpl_pair) == 2:
304                                        if special_entry := tmpl_pair[1].get('SpecialToken', {}).get('id'):
305                                            if special_entry in (special_sep, special_eos):
306                                                add_sep = True
307                                            if special_entry not in (special_sep, special_eos):
308                                                logger.warning(f'Unknown second separator token {special_entry!r} in TemplateProcessing<pair>')
309                                        else:
310                                            logger.warning(f'Unknown second middle sequence {tmpl_pair[1]!r} in TemplateProcessing<pair>')
311                                    self.add_special_token['sep'] = add_sep
312                                    if add_sep and not special_sep and tokenizer_config:
313                                        tokenizer_config['sep_token'] = special_eos
314                        continue
315        if not tokenizer_config:
316            return True
317        chat_template_alt = None
318        chat_template_json = path / 'chat_template.json'
319        chat_template_jinja = path / 'chat_template.jinja'
320        if chat_template_jinja.is_file():
321            with open(chat_template_jinja, encoding = 'utf-8') as f:
322                chat_template_alt = f.read()
323            if additional_templates := list((path / 'additional_chat_templates').glob('*.jinja')):
324                chat_template_alt = [{'name': 'default', 'template': chat_template_alt}]
325                for template_path in additional_templates:
326                    with open(template_path, encoding = 'utf-8') as fp:
327                        chat_template_alt.append({'name': template_path.stem, 'template': fp.read()})
328        elif chat_template_json.is_file():
329            with open(chat_template_json, encoding = 'utf-8') as f:
330                chat_template_alt = json.load(f).get('chat_template')
331        chat_template = tokenizer_config.get('chat_template', chat_template_alt)
332        if chat_template is None or isinstance(chat_template, (str, list)):
333            self.chat_template = chat_template
334        else:
335            logger.warning(f'Bad type for chat_template field in {tokenizer_config_file!r} - ignoring')
336        for typ in self.special_token_types:
337            add_entry = tokenizer_config.get(f'add_{typ}_token')
338            if isinstance(add_entry, bool):
339                self.add_special_token[typ] = add_entry
340            entry = tokenizer_config.get(f'{typ}_token')
341            if isinstance(entry, str):
342                tc_content = entry
343            elif isinstance(entry, dict):
344                entry_content = entry.get('content')
345                if not isinstance(entry_content, str):
346                    continue
347                tc_content = entry_content
348            else:
349                continue
350            # We only need the first match here.
351            maybe_token_id = next(
352                (atok.get('id') for atok in added_tokens if atok.get('content') == tc_content),
353                None,
354            )
355            self._set_special_token(typ, maybe_token_id)
356        return True
357
358    def _try_load_from_config_json(self, path: Path) -> bool:
359        config_file = path / 'config.json'
360        if not config_file.is_file():
361            return False
362        with open(config_file, encoding = 'utf-8') as f:
363            config = json.load(f)
364        for typ in self.special_token_types:
365            token_id = config.get(f'{typ}_token_id')
366            # If not found at root, check in text_config (for multimodal models like Kimi-VL)
367            if token_id is None and 'text_config' in config:
368                token_id = config['text_config'].get(f'{typ}_token_id')
369            self._set_special_token(typ, token_id)
370        return True
371
372
373@runtime_checkable
374class BaseVocab(Protocol):
375    tokenizer_model: ClassVar[str]
376    name: ClassVar[str]
377
378
379@runtime_checkable
380class Vocab(BaseVocab, Protocol):
381    vocab_size: int
382    added_tokens_dict: dict[str, int]
383    added_tokens_list: list[str]
384    fname_tokenizer: Path
385
386    def __init__(self, base_path: Path): ...
387    def all_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]: ...
388
389
390class NoVocab(BaseVocab):
391    tokenizer_model = "no_vocab"
392    name = "no_vocab"
393
394    def __repr__(self) -> str:
395        return "<NoVocab for a model without integrated vocabulary>"
396
397
398class BpeVocab(Vocab):
399    tokenizer_model = "gpt2"
400    name = "bpe"
401
402    def __init__(self, base_path: Path):
403        added_tokens: dict[str, int] = {}
404
405        if (fname_tokenizer := base_path / 'vocab.json').exists():
406            # "slow" tokenizer
407            with open(fname_tokenizer, encoding="utf-8") as f:
408                self.vocab = json.load(f)
409
410            try:
411                # FIXME: Verify that added tokens here _cannot_ overlap with the main vocab.
412                with open(base_path / 'added_tokens.json', encoding="utf-8") as f:
413                    added_tokens = json.load(f)
414            except FileNotFoundError:
415                pass
416        else:
417            # "fast" tokenizer
418            fname_tokenizer = base_path / 'tokenizer.json'
419
420            # if this fails, FileNotFoundError propagates to caller
421            with open(fname_tokenizer, encoding="utf-8") as f:
422                tokenizer_json = json.load(f)
423
424            tokenizer_model: dict[str, Any] = tokenizer_json['model']
425            if (
426                tokenizer_model['type'] != 'BPE' or tokenizer_model.get('byte_fallback', False)
427                or tokenizer_json['decoder']['type'] != 'ByteLevel'
428            ):
429                raise FileNotFoundError('Cannot find GPT-2 BPE tokenizer')
430
431            self.vocab = tokenizer_model["vocab"]
432
433            if (added := tokenizer_json.get('added_tokens')) is not None:
434                # Added tokens here can be duplicates of the main vocabulary.
435                added_tokens = {item['content']: item['id']
436                                for item in added
437                                if item['content'] not in self.vocab}
438
439        vocab_size   = len(self.vocab)
440        expected_ids = list(range(vocab_size, vocab_size + len(added_tokens)))
441        actual_ids   = sorted(added_tokens.values())
442        if expected_ids != actual_ids:
443            expected_end_id = vocab_size + len(actual_ids) - 1
444            raise ValueError(f"Expected the {len(actual_ids)} added token ID(s) to be sequential in the range "
445                             f"{vocab_size} - {expected_end_id}; got {actual_ids}")
446
447        items = sorted(added_tokens.items(), key=lambda text_idx: text_idx[1])
448        self.added_tokens_dict    = added_tokens
449        self.added_tokens_list    = [text for (text, idx) in items]
450        self.vocab_size_base      = vocab_size
451        self.vocab_size           = self.vocab_size_base + len(self.added_tokens_list)
452        self.fname_tokenizer      = fname_tokenizer
453
454    def bpe_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]:
455        reverse_vocab = {id: encoded_tok for encoded_tok, id in self.vocab.items()}
456
457        for i, _ in enumerate(self.vocab):
458            yield reverse_vocab[i], 0.0, gguf.TokenType.NORMAL
459
460    def added_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]:
461        for text in self.added_tokens_list:
462            score = -1000.0
463            yield text.encode("utf-8"), score, gguf.TokenType.CONTROL
464
465    def all_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]:
466        yield from self.bpe_tokens()
467        yield from self.added_tokens()
468
469    def __repr__(self) -> str:
470        return f"<BpeVocab with {self.vocab_size_base} base tokens and {len(self.added_tokens_list)} added tokens>"
471
472
473class SentencePieceVocab(Vocab):
474    tokenizer_model = "llama"
475    name = "spm"
476
477    def __init__(self, base_path: Path):
478        if SentencePieceProcessor is None:
479            raise RuntimeError("sentencepiece is not installed")
480
481        added_tokens: dict[str, int] = {}
482        if (fname_tokenizer := base_path / 'tokenizer.model').exists():
483            # normal location
484            try:
485                with open(base_path / 'added_tokens.json', encoding="utf-8") as f:
486                    added_tokens = json.load(f)
487            except FileNotFoundError:
488                pass
489        elif not (fname_tokenizer := base_path.parent / 'tokenizer.model').exists():
490            # not found in alternate location either
491            raise FileNotFoundError('Cannot find tokenizer.model')
492
493        self.sentencepiece_tokenizer = SentencePieceProcessor()
494        self.sentencepiece_tokenizer.LoadFromFile(str(fname_tokenizer))
495        vocab_size = self.sentencepiece_tokenizer.vocab_size()
496
497        new_tokens       = {id: piece for piece, id in added_tokens.items() if id >= vocab_size}
498        expected_new_ids = list(range(vocab_size, vocab_size + len(new_tokens)))
499        actual_new_ids   = sorted(new_tokens.keys())
500
501        if expected_new_ids != actual_new_ids:
502            raise ValueError(f"Expected new token IDs {expected_new_ids} to be sequential; got {actual_new_ids}")
503
504        # Token pieces that were added to the base vocabulary.
505        self.added_tokens_dict  = added_tokens
506        self.added_tokens_list  = [new_tokens[id] for id in actual_new_ids]
507        self.vocab_size_base    = vocab_size
508        self.vocab_size         = self.vocab_size_base + len(self.added_tokens_list)
509        self.fname_tokenizer    = fname_tokenizer
510
511    def sentencepiece_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]:
512        tokenizer = self.sentencepiece_tokenizer
513        for i in range(tokenizer.vocab_size()):
514            piece = tokenizer.IdToPiece(i)
515            text         = piece.encode("utf-8")
516            score: float = tokenizer.GetScore(i)
517
518            toktype = gguf.TokenType.NORMAL
519            if tokenizer.IsUnknown(i):
520                toktype = gguf.TokenType.UNKNOWN
521            if tokenizer.IsControl(i):
522                toktype = gguf.TokenType.CONTROL
523
524            # NOTE: I think added_tokens are user defined.
525            # ref: https://github.com/google/sentencepiece/blob/master/src/sentencepiece_model.proto
526            # if tokenizer.is_user_defined(i): toktype = gguf.TokenType.USER_DEFINED
527
528            if tokenizer.IsUnused(i):
529                toktype = gguf.TokenType.UNUSED
530            if tokenizer.IsByte(i):
531                toktype = gguf.TokenType.BYTE
532
533            yield text, score, toktype
534
535    def added_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]:
536        for text in self.added_tokens_list:
537            score = -1000.0
538            yield text.encode("utf-8"), score, gguf.TokenType.USER_DEFINED
539
540    def all_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]:
541        yield from self.sentencepiece_tokens()
542        yield from self.added_tokens()
543
544    def __repr__(self) -> str:
545        return f"<SentencePieceVocab with {self.vocab_size_base} base tokens and {len(self.added_tokens_list)} added tokens>"
546
547
548class LlamaHfVocab(Vocab):
549    tokenizer_model = "llama"
550    name = "hfft"
551
552    def __init__(self, base_path: Path):
553        fname_tokenizer = base_path / 'tokenizer.json'
554        # if this fails, FileNotFoundError propagates to caller
555        with open(fname_tokenizer, encoding='utf-8') as f:
556            tokenizer_json = json.load(f)
557
558        # pre-check so we know if we need transformers
559        tokenizer_model: dict[str, Any] = tokenizer_json['model']
560        is_llama3 = (
561            tokenizer_model['type'] == 'BPE' and tokenizer_model.get('ignore_merges', False)
562            and not tokenizer_model.get('byte_fallback', True)
563        )
564        if is_llama3:
565            raise TypeError('Llama 3 must be converted with BpeVocab')
566
567        if not is_llama3 and (
568            tokenizer_model['type'] != 'BPE' or not tokenizer_model.get('byte_fallback', False)
569            or tokenizer_json['decoder']['type'] != 'Sequence'
570        ):
571            raise FileNotFoundError('Cannot find Llama BPE tokenizer')
572
573        try:
574            from transformers import AutoTokenizer
575        except ImportError as e:
576            raise ImportError(
577                "To use LlamaHfVocab, please install the `transformers` package. "
578                "You can install it with `pip install transformers`."
579            ) from e
580
581        # Allow the tokenizer to default to slow or fast versions.
582        # Explicitly set tokenizer to use local paths.
583        self.tokenizer = AutoTokenizer.from_pretrained(
584            base_path,
585            cache_dir=base_path,
586            local_files_only=True,
587        )
588        assert self.tokenizer.is_fast  # assume tokenizer.json is used  # ty: ignore[unresolved-attribute]
589
590        # Initialize lists and dictionaries for added tokens
591        self.added_tokens_list = []
592        self.added_tokens_dict = dict()
593        self.added_tokens_ids  = set()
594
595        # Process added tokens
596        for tok, tokidx in sorted(
597            self.tokenizer.get_added_vocab().items(), key=lambda x: x[1]  # ty: ignore[unresolved-attribute]
598        ):
599            # Only consider added tokens that are not in the base vocabulary
600            if tokidx >= self.tokenizer.vocab_size:  # ty: ignore[unresolved-attribute]
601                self.added_tokens_list.append(tok)
602                self.added_tokens_dict[tok] = tokidx
603                self.added_tokens_ids.add(tokidx)
604
605        # Store special tokens and their IDs
606        self.specials = {
607            tok: self.tokenizer.get_vocab()[tok]  # ty: ignore[unresolved-attribute]
608            for tok in self.tokenizer.all_special_tokens  # ty: ignore[unresolved-attribute]
609        }
610        self.special_ids = set(self.tokenizer.all_special_ids)  # ty: ignore[unresolved-attribute]
611
612        # Set vocabulary sizes
613        self.vocab_size_base = self.tokenizer.vocab_size  # ty: ignore[unresolved-attribute]
614        self.vocab_size      = self.vocab_size_base + len(self.added_tokens_list)
615
616        self.fname_tokenizer = fname_tokenizer
617
618    def hf_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]:
619        reverse_vocab = {
620            id: encoded_tok for encoded_tok, id in self.tokenizer.get_vocab().items()  # ty: ignore[unresolved-attribute]
621        }
622
623        for token_id in range(self.vocab_size_base):
624            # Skip processing added tokens here
625            if token_id in self.added_tokens_ids:
626                continue
627
628            # Convert token text to bytes
629            token_text = reverse_vocab[token_id].encode("utf-8")
630
631            # Yield token text, score, and type
632            yield token_text, self.get_token_score(token_id), self.get_token_type(
633                token_id, token_text, self.special_ids  # Reuse already stored special IDs
634            )
635
636    def get_token_type(self, token_id: int, token_text: bytes, special_ids: set[int]) -> gguf.TokenType:
637        # Special case for byte tokens
638        if re.fullmatch(br"<0x[0-9A-Fa-f]{2}>", token_text):
639            return gguf.TokenType.BYTE
640
641        # Determine token type based on whether it's a special token
642        return gguf.TokenType.CONTROL if token_id in special_ids else gguf.TokenType.NORMAL
643
644    def get_token_score(self, token_id: int) -> float:
645        # Placeholder for actual logic to determine the token's score
646        # This needs to be implemented based on specific requirements
647        return -1000.0  # Default score
648
649    def added_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]:
650        for text in self.added_tokens_list:
651            if text in self.specials:
652                toktype = self.get_token_type(self.specials[text], b'', self.special_ids)
653                score = self.get_token_score(self.specials[text])
654            else:
655                toktype = gguf.TokenType.USER_DEFINED
656                score = -1000.0
657
658            yield text.encode("utf-8"), score, toktype
659
660    def has_newline_token(self):
661        return "<0x0A>" in self.tokenizer.vocab or "\n" in self.tokenizer.vocab  # ty: ignore[unresolved-attribute]
662
663    def all_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]:
664        yield from self.hf_tokens()
665        yield from self.added_tokens()
666
667    def __repr__(self) -> str:
668        return f"<LlamaHfVocab with {self.vocab_size_base} base tokens and {len(self.added_tokens_list)} added tokens>"
669
670
671class MistralTokenizerType(str, Enum):
672    spm = "spm"
673    tekken = "tekken"
674
675
676# Copied from Transformers (Apache 2.0)
677# https://github.com/huggingface/transformers/blob/main/src/transformers/convert_slow_tokenizer.py#L1544
678
679def bytes_to_unicode() -> dict[int, str]:
680    """
681    Returns list of utf-8 byte and a mapping to unicode strings. We specifically avoids mapping to whitespace/control
682    characters the bpe code barfs on.
683
684    The reversible bpe codes work on unicode strings. This means you need a large # of unicode characters in your vocab
685    if you want to avoid UNKs. When you're at something like a 10B token dataset you end up needing around 5K for
686    decent coverage. This is a significant percentage of your normal, say, 32K bpe vocab. To avoid that, we want lookup
687    tables between utf-8 bytes and unicode strings.
688    """
689    bs = (
690        list(range(ord("!"), ord("~") + 1))
691        + list(range(ord("¡"), ord("¬") + 1))
692        + list(range(ord("®"), ord("ÿ") + 1))
693    )
694    cs = bs[:]
695    n = 0
696    for b in range(2**8):
697        if b not in bs:
698            bs.append(b)
699            cs.append(2**8 + n)
700            n += 1
701    cs_str = [chr(n) for n in cs]
702    return dict(zip(bs, cs_str))
703
704
705class MistralVocab(Vocab):
706    tokenizer_model = "mistral"
707    name = "mistral"
708
709    added_tokens_dict: dict[str, int] = {}
710    added_tokens_list: list[str] = []
711
712    def __init__(self, base_path: Path):
713        if not _mistral_common_installed:
714            raise ImportError(
715                "To use MistralVocab, please install the `mistral-common` package. "
716                "You can install it with `pip install mistral-common`."
717            )
718        assert _filter_valid_tokenizer_files is not None, "mistral_common is not installed"
719        assert MistralTokenizer is not None, "mistral_common is not installed"
720        assert Tekkenizer is not None, "mistral_common is not installed"
721
722        logger.info(f"Loading Mistral tokenizer from {base_path}")
723
724        # Find the tokenizer files
725        all_files = [f.as_posix() for f in base_path.glob("**/*") if f.is_file()]
726
727        if get_one_valid_tokenizer_file is not None:
728            tokenizer_file_path = get_one_valid_tokenizer_file(all_files)
729        else:
730            valid_tokenizer_files = _filter_valid_tokenizer_files(all_files)
731
732            if len(valid_tokenizer_files) == 0:
733                raise ValueError(f"No tokenizer file found in the directory: {base_path}")
734            # If there are multiple tokenizer files, we use tekken.json if it exists, otherwise the versioned one.
735            if len(valid_tokenizer_files) > 1:
736                if "tekken.json" in valid_tokenizer_files:
737                    tokenizer_file = "tekken.json"
738                else:
739                    tokenizer_file = sorted(valid_tokenizer_files)[-1]
740                logger.warning(
741                    f"Multiple tokenizer files found in {base_path}. Using {tokenizer_file}"
742                )
743            else:
744                tokenizer_file = valid_tokenizer_files[0]
745
746            tokenizer_file_path = base_path / tokenizer_file
747
748        self.tokenizer: Any = MistralTokenizer.from_file(
749            tokenizer_file_path
750        ).instruct_tokenizer.tokenizer
751        self.tokenizer_type = (
752            MistralTokenizerType.tekken
753            if isinstance(self.tokenizer, Tekkenizer)
754            else MistralTokenizerType.spm
755        )
756        self.vocab_size = self.tokenizer.n_words
757        self.fname_tokenizer = tokenizer_file_path
758        self._name = (
759            "mistral-" + self.tokenizer_type.value + "-" + self.tokenizer.version
760        )
761
762    @property
763    def tokenizer_name(self) -> str:
764        return self._name
765
766    @property
767    def gguf_tokenizer_model(self) -> str:
768        return "llama" if self.tokenizer_type == MistralTokenizerType.spm else "gpt2"
769
770    def _sentencepiece_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]:
771        assert SentencePieceTokenizer is not None, "mistral_common is not installed"
772        assert isinstance(self.tokenizer, SentencePieceTokenizer), (
773            f"Expected SentencePieceTokenizer, got {type(self.tokenizer)}"
774        )
775
776        for i in range(self.tokenizer._model.vocab_size()):
777            piece = self.tokenizer._model.IdToPiece(i)
778            text = piece.encode("utf-8")
779            score: float = self.tokenizer._model.GetScore(i)
780
781            toktype = gguf.TokenType.NORMAL
782            if self.tokenizer._model.IsUnknown(i):
783                toktype = gguf.TokenType.UNKNOWN
784            if self.tokenizer._model.IsControl(i):
785                toktype = gguf.TokenType.CONTROL
786
787            if self.tokenizer._model.IsUnused(i):
788                toktype = gguf.TokenType.UNUSED
789            if self.tokenizer._model.IsByte(i):
790                toktype = gguf.TokenType.BYTE
791
792            yield text, score, toktype
793
794    def _tekken_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]:
795        assert Tekkenizer is not None, "mistral_common is not installed"
796        assert isinstance(self.tokenizer, Tekkenizer), (
797            f"Expected Tekkenizer, got {type(self.tokenizer)}"
798        )
799
800        byte_encoder = bytes_to_unicode()
801        for token_id in range(self.tokenizer.num_special_tokens):
802            yield (
803                self.tokenizer.id_to_piece(token_id).encode("utf-8"),
804                0,
805                gguf.TokenType.CONTROL
806            )
807        for token in self.tokenizer._tekken_token2id_nospecial:
808            yield (
809                self.token_bytes_to_string(token, byte_encoder).encode("utf-8"),
810                0,
811                gguf.TokenType.NORMAL,
812            )
813
814    def get_token_id(self, token: str) -> int:
815        assert SentencePieceTokenizer is not None and Tekkenizer is not None, "mistral_common is not installed"
816        if self.tokenizer_type == MistralTokenizerType.spm:
817            assert isinstance(self.tokenizer, SentencePieceTokenizer)
818            return self.tokenizer._vocab.index(token)
819        elif self.tokenizer_type == MistralTokenizerType.tekken:
820            assert isinstance(self.tokenizer, Tekkenizer)
821            return (
822                self.tokenizer._vocab.index(token) + self.tokenizer.num_special_tokens
823            )
824        else:
825            raise ValueError(f"Unknown tokenizer type: {self.tokenizer_type}")
826
827    @property
828    def bos_id(self) -> int:
829        return self.tokenizer.bos_id
830
831    @property
832    def eos_id(self) -> int:
833        return self.tokenizer.eos_id
834
835    @property
836    def pad_id(self) -> int:
837        if self.tokenizer.pad_id == -1:
838            return self.eos_id
839        return self.tokenizer.pad_id
840
841    @property
842    def unk_id(self) -> int:
843        return self.tokenizer.unk_id
844
845    @property
846    def bos_token(self) -> str:
847        return self.tokenizer.id_to_piece(self.tokenizer.bos_id)
848
849    @property
850    def eos_token(self) -> str:
851        return self.tokenizer.id_to_piece(self.tokenizer.eos_id)
852
853    @property
854    def pad_token(self) -> str:
855        return self.tokenizer.id_to_piece(self.tokenizer.pad_id)
856
857    @property
858    def unk_token(self) -> str:
859        return self.tokenizer.id_to_piece(self.tokenizer.unk_id)
860
861    def all_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]:
862        if self.tokenizer_type == MistralTokenizerType.spm:
863            yield from self._sentencepiece_tokens()
864
865        elif self.tokenizer_type == MistralTokenizerType.tekken:
866            yield from self._tekken_tokens()
867
868        else:
869            raise ValueError(f"Unknown tokenizer type: {self.tokenizer_type}")
870
871    @staticmethod
872    def token_bytes_to_string(b, byte_encoder):
873        return "".join([byte_encoder[ord(char)] for char in b.decode("latin-1")])
874
875    def extract_vocab_merges_from_model(self):
876        # Adapted from Transformers (Apache 2.0)
877        # https://github.com/huggingface/transformers/blob/main/src/transformers/convert_slow_tokenizer.py
878        assert Tekkenizer is not None and isinstance(self.tokenizer, Tekkenizer), (
879            f"Expected Tekkenizer, got {type(self.tokenizer)}"
880        )
881        mergeable_ranks = self.tokenizer._model._mergeable_ranks
882        token_bytes_map = {
883            rank: token_bytes for token_bytes, rank in mergeable_ranks.items()
884        }
885        merge_pairs = []
886
887        # Sort vocab by rank to ensure correct merge order
888        for i in range(256, self.vocab_size - self.tokenizer.num_special_tokens):
889            merged_token = token_bytes_map[i]
890            local = []
891            for j in range(1, len(merged_token)):
892                left = merged_token[:j]
893                right = merged_token[j:]
894                if (
895                    left in mergeable_ranks
896                    and right in mergeable_ranks
897                    and (left + right) in mergeable_ranks
898                ):
899                    local.append((left, right, i))
900            if not local:
901                raise ValueError(
902                    f"Could not find valid merge for token at rank {i}: {merged_token.decode('latin-1')}"
903                )
904            local = sorted(
905                local,
906                key=lambda x: (mergeable_ranks[x[0]], mergeable_ranks[x[1]]),
907                reverse=False,
908            )
909            merge_pairs.extend(local)
910        merge_pairs = sorted(merge_pairs, key=lambda val: val[2], reverse=False)
911
912        byte_encoder = bytes_to_unicode()
913
914        decoded_merge_pairs = [
915            [
916                self.token_bytes_to_string(val[0], byte_encoder),
917                self.token_bytes_to_string(val[1], byte_encoder),
918            ]
919            for val in merge_pairs
920        ]
921
922        merges = [
923            " ".join(
924                [
925                    # ensure the spaces are properly encoded
926                    "".join(chr(ord(c) + 256) if c == " " else c for c in part)
927                    for part in pair
928                ]
929            )
930            for pair in decoded_merge_pairs
931        ]
932
933        return merges