| from transformers import PreTrainedTokenizer |
| from typing import Dict, List, Optional, Any |
|
|
|
|
| class ByteTokenizer(PreTrainedTokenizer): |
| model_input_names = ["input_ids", "attention_mask"] |
|
|
| def __init__( |
| self, |
| bos_token="<|bos|>", |
| eos_token="<|eos|>", |
| pad_token="<|pad|>", |
| vocab_size=384, |
| **kwargs, |
| ): |
| self.pad_idx = 0 |
| self.bos_idx = 254 |
| self.eos_idx = 255 |
| self._vocab_size = vocab_size |
|
|
| self.byte_to_token = [f"<byte_{i}>" for i in range(256)] |
| self.token_to_byte = {t: i for i, t in enumerate(self.byte_to_token)} |
|
|
| super().__init__( |
| bos_token=bos_token, |
| eos_token=eos_token, |
| pad_token=pad_token, |
| **kwargs, |
| ) |
|
|
| @property |
| def vocab_size(self) -> int: |
| return self._vocab_size |
|
|
| def get_vocab(self) -> Dict[str, int]: |
| vocab = {t: i for i, t in enumerate(self.byte_to_token)} |
| vocab.update( |
| { |
| self.bos_token: self.bos_idx, |
| self.eos_token: self.eos_idx, |
| self.pad_token: self.pad_idx, |
| } |
| ) |
| return vocab |
|
|
| def _tokenize(self, text, **kwargs): |
| return [self.byte_to_token[b] for b in text.encode("utf-8")] |
|
|
| def _convert_token_to_id(self, token): |
| if token == self.bos_token: |
| return self.bos_idx |
| if token == self.eos_token: |
| return self.eos_idx |
| if token == self.pad_token: |
| return self.pad_idx |
| return self.token_to_byte.get(token, self.pad_idx) |
|
|
| def _convert_id_to_token(self, index): |
| if index == self.bos_idx: |
| return self.bos_token |
| if index == self.eos_idx: |
| return self.eos_token |
| if index == self.pad_idx: |
| return self.pad_token |
| if 0 <= index < 256: |
| return self.byte_to_token[index] |
| return f"<unk_{index}>" |
|
|
| def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None): |
| return [self.bos_idx] + token_ids_0 + [self.eos_idx] |
|
|
| def _decode( |
| self, token_ids: List[int], skip_special_tokens: bool = False, **kwargs |
| ) -> str: |
| clean_ids = [] |
| for i in token_ids: |
| if skip_special_tokens and i in [self.bos_idx, self.eos_idx, self.pad_idx]: |
| continue |
| if 0 <= i < 256: |
| clean_ids.append(i) |
| return bytes(clean_ids).decode("utf-8", errors="ignore") |
|
|
| def save_vocabulary( |
| self, save_directory: str, filename_prefix: Optional[str] = None |
| ) -> tuple: |
| return () |
|
|
|
|
| __all__ = [ |
| "ByteTokenizer", |
| ] |
|
|
