| from __future__ import annotations |
|
|
| import abc |
| from typing import ( |
| List, |
| Optional, |
| Any, |
| ) |
|
|
| import llama_cpp |
| from llama_cpp.llama_types import List |
|
|
|
|
| class BaseLlamaTokenizer(abc.ABC): |
| @abc.abstractmethod |
| def tokenize( |
| self, text: bytes, add_bos: bool = True, special: bool = True |
| ) -> List[int]: |
| """Tokenize the text into tokens. |
| |
| Args: |
| text: The utf-8 encoded string to tokenize. |
| add_bos: Whether to add a beginning of sequence token. |
| special: Whether to tokenize special tokens. |
| """ |
| raise NotImplementedError |
|
|
| @abc.abstractmethod |
| def detokenize( |
| self, |
| tokens: List[int], |
| prev_tokens: Optional[List[int]] = None, |
| special: bool = False, |
| ) -> bytes: |
| """Detokenize the tokens into text. |
| |
| Args: |
| tokens: The list of tokens to detokenize. |
| prev_tokens: The list of previous tokens. Offset mapping will be performed if provided. |
| special: Whether to detokenize special tokens. |
| """ |
| raise NotImplementedError |
|
|
|
|
| class LlamaTokenizer(BaseLlamaTokenizer): |
| def __init__(self, llama: llama_cpp.Llama): |
| self._model = llama._model |
|
|
| def tokenize( |
| self, text: bytes, add_bos: bool = True, special: bool = True |
| ) -> List[int]: |
| return self._model.tokenize(text, add_bos=add_bos, special=special) |
|
|
| def detokenize( |
| self, |
| tokens: List[int], |
| prev_tokens: Optional[List[int]] = None, |
| special: bool = False, |
| ) -> bytes: |
| return self._model.detokenize(tokens, special=special) |
|
|
| def encode( |
| self, text: str, add_bos: bool = True, special: bool = True |
| ) -> List[int]: |
| return self.tokenize( |
| text.encode("utf-8", errors="ignore"), add_bos=add_bos, special=special |
| ) |
|
|
| def decode(self, tokens: List[int]) -> str: |
| return self.detokenize(tokens).decode("utf-8", errors="ignore") |
|
|
| @classmethod |
| def from_ggml_file(cls, path: str) -> "LlamaTokenizer": |
| return cls(llama_cpp.Llama(model_path=path, vocab_only=True)) |
|
|
|
|
| class LlamaHFTokenizer(BaseLlamaTokenizer): |
| def __init__(self, hf_tokenizer: Any): |
| self.hf_tokenizer = hf_tokenizer |
|
|
| def tokenize( |
| self, text: bytes, add_bos: bool = True, special: bool = True |
| ) -> List[int]: |
| return self.hf_tokenizer.encode( |
| text.decode("utf-8", errors="ignore"), add_special_tokens=special |
| ) |
|
|
| def detokenize( |
| self, |
| tokens: List[int], |
| prev_tokens: Optional[List[int]] = None, |
| special: bool = False, |
| ) -> bytes: |
| skip_special_tokens = not special |
| if prev_tokens is not None: |
| text = self.hf_tokenizer.decode( |
| prev_tokens + tokens, skip_special_tokens=skip_special_tokens |
| ).encode("utf-8", errors="ignore") |
| prev_text = self.hf_tokenizer.decode( |
| prev_tokens, skip_special_tokens=skip_special_tokens |
| ).encode("utf-8", errors="ignore") |
| return text[len(prev_text) :] |
| else: |
| return self.hf_tokenizer.decode( |
| tokens, skip_special_tokens=skip_special_tokens |
| ).encode("utf-8", errors="ignore") |
|
|
| @classmethod |
| def from_pretrained(cls, pretrained_model_name_or_path: str) -> "LlamaHFTokenizer": |
| try: |
| from transformers import AutoTokenizer |
| except ImportError: |
| raise ImportError( |
| "The `transformers` library is required to use the `HFTokenizer`." |
| "You can install it with `pip install transformers`." |
| ) |
| hf_tokenizer = AutoTokenizer.from_pretrained( |
| pretrained_model_name_or_path=pretrained_model_name_or_path |
| ) |
| return cls(hf_tokenizer) |
|
|
