|
|
|
|
|
from __future__ import annotations
|
|
|
|
|
|
import torch
|
|
|
|
|
|
import numpy as np
|
|
|
|
|
|
from os import PathLike
|
|
|
from typing import List, Tuple
|
|
|
|
|
|
from tokenizers import Tokenizer
|
|
|
from transformers.tokenization_utils import PreTrainedTokenizer
|
|
|
from transformers.tokenization_utils_base import BatchEncoding, TruncationStrategy
|
|
|
from transformers.utils.generic import TensorType, PaddingStrategy
|
|
|
|
|
|
|
|
|
EMPTY: str = ""
|
|
|
|
|
|
|
|
|
class ByteTokenizer(PreTrainedTokenizer):
|
|
|
|
|
|
"""UTF-8 Encoder."""
|
|
|
|
|
|
@classmethod
|
|
|
def from_pretrained(cls, model_id: str | PathLike, **kwargs) -> ByteTokenizer:
|
|
|
|
|
|
return cls(**kwargs, byte_level=True)
|
|
|
|
|
|
@property
|
|
|
def vocab_size(self) -> int:
|
|
|
|
|
|
return 512
|
|
|
|
|
|
@property
|
|
|
def byte_level(self) -> bool:
|
|
|
|
|
|
return self.init_kwargs.get('byte_level', True)
|
|
|
|
|
|
def get_vocab(self) -> Dict[str, int]:
|
|
|
|
|
|
return {chr(i): i for i in range(self.vocab_size)}
|
|
|
|
|
|
def __len__(self) -> int:
|
|
|
|
|
|
return self.vocab_size
|
|
|
|
|
|
def clamp(self, n: int) -> int:
|
|
|
|
|
|
return max(32, min(n, self.vocab_size))
|
|
|
|
|
|
def _tokenize(self, text: str, **kwargs) -> List[str]:
|
|
|
|
|
|
return list(text)
|
|
|
|
|
|
def byte_tokenize(self, text: str) -> np.ndarray:
|
|
|
|
|
|
return np.frombuffer(text.encode('utf-8'), dtype=np.uint8)
|
|
|
|
|
|
def _convert_token_to_id(self, token: str) -> int:
|
|
|
|
|
|
return self.clamp(ord(token))
|
|
|
|
|
|
def _convert_id_to_token(self, index: int) -> str:
|
|
|
|
|
|
return chr(self.clamp(index))
|
|
|
|
|
|
def convert_tokens_to_string(self, tokens: List[str]) -> str:
|
|
|
|
|
|
return EMPTY.join(tokens)
|
|
|
|
|
|
def _decode(self, token_ids: List[int], **kwargs) -> str:
|
|
|
|
|
|
indices = np.asarray(token_ids, dtype=np.uint8)
|
|
|
|
|
|
return (
|
|
|
indices.clip(min=32, max=self.vocab_size, out=indices)
|
|
|
.tobytes()
|
|
|
.decode('utf-8')
|
|
|
)
|
|
|
|
|
|
def _encode_plus(self, text: str, **kwargs) -> BatchEncoding:
|
|
|
|
|
|
first_ids = self.byte_tokenize(text).tolist()
|
|
|
|
|
|
return self.prepare_for_model(
|
|
|
first_ids,
|
|
|
pair_ids=None,
|
|
|
add_special_tokens=kwargs.get('add_special_tokens', False),
|
|
|
padding=kwargs.get('padding_strategy', PaddingStrategy.DO_NOT_PAD).value,
|
|
|
truncation=kwargs.get('truncation_strategy', TruncationStrategy.DO_NOT_TRUNCATE).value,
|
|
|
max_length=kwargs.get('max_length'),
|
|
|
stride=kwargs.get('stride', 0),
|
|
|
pad_to_multiple_of=kwargs.get('pad_to_multiple_of'),
|
|
|
return_tensors=kwargs.get('return_tensors'),
|
|
|
prepend_batch_axis=True,
|
|
|
return_attention_mask=kwargs.get('return_attention_mask'),
|
|
|
return_token_type_ids=kwargs.get('return_token_type_ids'),
|
|
|
return_overflowing_tokens=kwargs.get('return_overflowing_tokens', False),
|
|
|
return_special_tokens_mask=kwargs.get('return_special_tokens_mask', False),
|
|
|
return_length=kwargs.get('return_length', False),
|
|
|
verbose=kwargs.get('verbose', True),
|
|
|
)
|
|
|
|
|
|
def _batch_encode_plus(self, batch_text_or_text_pairs: List[str], **kwargs) -> BatchEncoding:
|
|
|
|
|
|
input_ids = [(self.byte_tokenize(text).tolist(), None) for text in batch_text_or_text_pairs]
|
|
|
|
|
|
return self._batch_prepare_for_model(
|
|
|
input_ids,
|
|
|
add_special_tokens=kwargs.get('add_special_tokens', False),
|
|
|
padding_strategy=kwargs.get('padding_strategy', PaddingStrategy.DO_NOT_PAD),
|
|
|
truncation_strategy=kwargs.get('truncation_strategy', TruncationStrategy.DO_NOT_TRUNCATE),
|
|
|
max_length=kwargs.get('max_length'),
|
|
|
stride=kwargs.get('stride', 0),
|
|
|
pad_to_multiple_of=kwargs.get('pad_to_multiple_of'),
|
|
|
return_attention_mask=kwargs.get('return_attention_mask'),
|
|
|
return_token_type_ids=kwargs.get('return_token_type_ids'),
|
|
|
return_overflowing_tokens=kwargs.get('return_overflowing_tokens', False),
|
|
|
return_special_tokens_mask=kwargs.get('return_special_tokens_mask', False),
|
|
|
return_length=kwargs.get('return_length', False),
|
|
|
return_tensors=kwargs.get('return_tensors'),
|
|
|
verbose=kwargs.get('verbose', True),
|
|
|
)
|
|
|
|
|
|
def _save_pretrained(
|
|
|
self, save_directory: str | PathLike, file_names: Tuple[str], **kwargs
|
|
|
) -> Tuple[str]:
|
|
|
|
|
|
return file_names
|
|
|
|
