144 lines
4.5 KiB
Python
144 lines
4.5 KiB
Python
# Adapted from:
|
|
# https://github.com/outlines-dev/outlines/blob/0355ab4272a5d7e4d94c4a53a52593f885b81a61/outlines/models/tokenizer.py
|
|
# https://github.com/outlines-dev/outlines/blob/0355ab4272a5d7e4d94c4a53a52593f885b81a61/outlines/models/transformers.py
|
|
from abc import abstractmethod
|
|
from typing import Dict, Hashable, List, Protocol, Set, Tuple, Union
|
|
|
|
import numpy as np
|
|
import torch
|
|
from numpy.typing import NDArray
|
|
|
|
|
|
class Tokenizer(Protocol, Hashable):
|
|
eos_token: str
|
|
eos_token_id: int
|
|
pad_token_id: int
|
|
vocabulary: Dict[str, int]
|
|
special_tokens: Set[int]
|
|
|
|
@abstractmethod
|
|
def encode(
|
|
self, prompt: Union[str, List[str]]
|
|
) -> Tuple[NDArray[np.int64], NDArray[np.int64]]:
|
|
"""Translate the input prompts into NumPy arrays of token ids and attention mask."""
|
|
...
|
|
|
|
@abstractmethod
|
|
def decode(self, token_ids: NDArray[np.int64]) -> List[str]:
|
|
"""Translate an array of token ids to a string or list of strings."""
|
|
...
|
|
|
|
@abstractmethod
|
|
def convert_token_to_string(self, token: str) -> str:
|
|
"""Convert a token to its equivalent string.
|
|
|
|
This is for instance useful for BPE tokenizers where whitespaces are
|
|
represented by the special characted `Ġ`. This prevents matching a raw
|
|
token that includes `Ġ` with a string.
|
|
|
|
"""
|
|
...
|
|
|
|
|
|
def get_llama_tokenizer_types():
|
|
"""Get all the Llama tokenizer types/classes that need work-arounds.
|
|
|
|
When they can't be imported, a dummy class is created.
|
|
|
|
"""
|
|
try:
|
|
from transformers.models.llama import LlamaTokenizer
|
|
except ImportError:
|
|
|
|
class LlamaTokenizer: # type: ignore
|
|
pass
|
|
|
|
try:
|
|
from transformers.models.llama import LlamaTokenizerFast
|
|
except ImportError:
|
|
|
|
class LlamaTokenizerFast: # type: ignore
|
|
pass
|
|
|
|
try:
|
|
from transformers.models.code_llama import CodeLlamaTokenizer
|
|
except ImportError:
|
|
|
|
class CodeLlamaTokenizer: # type: ignore
|
|
pass
|
|
|
|
try:
|
|
from transformers.models.code_llama import CodeLlamaTokenizerFast
|
|
except ImportError:
|
|
|
|
class CodeLlamaTokenizerFast: # type: ignore
|
|
pass
|
|
|
|
return (
|
|
LlamaTokenizer,
|
|
LlamaTokenizerFast,
|
|
CodeLlamaTokenizer,
|
|
CodeLlamaTokenizerFast,
|
|
)
|
|
|
|
|
|
class TransformerTokenizer(Tokenizer):
|
|
"""Represents a tokenizer for models in the `transformers` library."""
|
|
|
|
def __init__(self, model_name: str, **kwargs):
|
|
from transformers import AutoTokenizer
|
|
|
|
kwargs.setdefault("padding_side", "left")
|
|
self.model_name = model_name
|
|
# TODO: Do something to make this hashable?
|
|
self.kwargs = kwargs
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_name, **kwargs)
|
|
self.eos_token_id = self.tokenizer.eos_token_id
|
|
self.eos_token = self.tokenizer.eos_token
|
|
|
|
if not self.tokenizer.pad_token_id:
|
|
self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
|
|
self.pad_token_id = self.eos_token_id
|
|
else:
|
|
self.pad_token_id = self.tokenizer.pad_token_id
|
|
self.pad_token = self.tokenizer.pad_token
|
|
|
|
self.special_tokens = set(self.tokenizer.all_special_tokens)
|
|
|
|
self.vocabulary = self.tokenizer.get_vocab()
|
|
self.is_llama = isinstance(self.tokenizer, get_llama_tokenizer_types())
|
|
|
|
def encode(
|
|
self, prompt: Union[str, List[str]], **kwargs
|
|
) -> Tuple[torch.LongTensor, torch.LongTensor]:
|
|
kwargs["padding"] = True
|
|
kwargs["return_tensors"] = "pt"
|
|
output = self.tokenizer(prompt, **kwargs)
|
|
return output["input_ids"], output["attention_mask"]
|
|
|
|
def decode(self, token_ids: torch.LongTensor) -> List[str]:
|
|
text = self.tokenizer.batch_decode(token_ids, skip_special_tokens=True)
|
|
return text
|
|
|
|
def convert_token_to_string(self, token: str) -> str:
|
|
from transformers.file_utils import SPIECE_UNDERLINE
|
|
|
|
string = self.tokenizer.convert_tokens_to_string([token])
|
|
|
|
if self.is_llama:
|
|
# A hack to handle missing spaces to HF's Llama tokenizers
|
|
if token.startswith(SPIECE_UNDERLINE) or token == "<0x20>":
|
|
return " " + string
|
|
|
|
return string
|
|
|
|
def __eq__(self, other):
|
|
if isinstance(other, type(self)):
|
|
return other.model_name == self.model_name and other.kwargs == self.kwargs
|
|
return NotImplemented
|
|
|
|
def __hash__(self):
|
|
from datasets.fingerprint import Hasher
|
|
|
|
return hash(Hasher.hash(self.tokenizer))
|