hymenjj/llama-cpp-python-prebuilt

1from __future__ import annotations2 3import abc4from typing import (5    List,6    Optional,7    Any,8)9 10import llama_cpp11from llama_cpp.llama_types import List12 13 14class BaseLlamaTokenizer(abc.ABC):15    @abc.abstractmethod16    def tokenize(17        self, text: bytes, add_bos: bool = True, special: bool = True18    ) -> List[int]:19        """Tokenize the text into tokens.20 21        Args:22            text: The utf-8 encoded string to tokenize.23            add_bos: Whether to add a beginning of sequence token.24            special: Whether to tokenize special tokens.25        """26        raise NotImplementedError27 28    @abc.abstractmethod29    def detokenize(30        self,31        tokens: List[int],32        prev_tokens: Optional[List[int]] = None,33        special: bool = False,34    ) -> bytes:35        """Detokenize the tokens into text.36 37        Args:38            tokens: The list of tokens to detokenize.39            prev_tokens: The list of previous tokens. Offset mapping will be performed if provided.40            special: Whether to detokenize special tokens.41        """42        raise NotImplementedError43 44 45class LlamaTokenizer(BaseLlamaTokenizer):46    def __init__(self, llama: llama_cpp.Llama):47        self._model = llama._model  # type: ignore48 49    def tokenize(50        self, text: bytes, add_bos: bool = True, special: bool = True51    ) -> List[int]:52        return self._model.tokenize(text, add_bos=add_bos, special=special)53 54    def detokenize(55        self,56        tokens: List[int],57        prev_tokens: Optional[List[int]] = None,58        special: bool = False,59    ) -> bytes:60        return self._model.detokenize(tokens, special=special)61 62    def encode(63        self, text: str, add_bos: bool = True, special: bool = True64    ) -> List[int]:65        return self.tokenize(66            text.encode("utf-8", errors="ignore"), add_bos=add_bos, special=special67        )68 69    def decode(self, tokens: List[int]) -> str:70        return self.detokenize(tokens).decode("utf-8", errors="ignore")71 72    @classmethod73    def from_ggml_file(cls, path: str) -> "LlamaTokenizer":74        return cls(llama_cpp.Llama(model_path=path, vocab_only=True))75 76 77class LlamaHFTokenizer(BaseLlamaTokenizer):78    def __init__(self, hf_tokenizer: Any):79        self.hf_tokenizer = hf_tokenizer80 81    def tokenize(82        self, text: bytes, add_bos: bool = True, special: bool = True83    ) -> List[int]:84        return self.hf_tokenizer.encode(85            text.decode("utf-8", errors="ignore"), add_special_tokens=special86        )87 88    def detokenize(89        self,90        tokens: List[int],91        prev_tokens: Optional[List[int]] = None,92        special: bool = False,93    ) -> bytes:94        skip_special_tokens = not special95        if prev_tokens is not None:96            text = self.hf_tokenizer.decode(97                prev_tokens + tokens, skip_special_tokens=skip_special_tokens98            ).encode("utf-8", errors="ignore")99            prev_text = self.hf_tokenizer.decode(100                prev_tokens, skip_special_tokens=skip_special_tokens101            ).encode("utf-8", errors="ignore")102            return text[len(prev_text) :]103        else:104            return self.hf_tokenizer.decode(105                tokens, skip_special_tokens=skip_special_tokens106            ).encode("utf-8", errors="ignore")107 108    @classmethod109    def from_pretrained(cls, pretrained_model_name_or_path: str) -> "LlamaHFTokenizer":110        try:111            from transformers import AutoTokenizer112        except ImportError:113            raise ImportError(114                "The `transformers` library is required to use the `HFTokenizer`."115                "You can install it with `pip install transformers`."116            )117        hf_tokenizer = AutoTokenizer.from_pretrained(118            pretrained_model_name_or_path=pretrained_model_name_or_path119        )120        return cls(hf_tokenizer)121