File size: 4,456 Bytes
da49047 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | import json
import os
from transformers import PreTrainedTokenizer
BYTE_PREFIX = "<0x"
PAD_TOKEN = "<pad>"
BOS_TOKEN = "<bos>"
EOS_TOKEN = "<eos>"
class BETByteTokenizer(PreTrainedTokenizer):
"""Lossless UTF-8 byte tokenizer used by BET.
IDs:
0..255 -> raw byte values
256 -> PAD
257 -> BOS
258 -> EOS
No UNK token is required because every UTF-8 string is representable as bytes.
"""
vocab_files_names = {"vocab_file": "byte_vocab.json"}
model_input_names = ["input_ids", "attention_mask"]
def __init__(
self,
vocab_file=None,
pad_token=PAD_TOKEN,
bos_token=BOS_TOKEN,
eos_token=EOS_TOKEN,
unk_token=None,
model_max_length=1024,
padding_side="left",
clean_up_tokenization_spaces=False,
**kwargs,
):
# Transformers v5 loads values from tokenizer_config.json into this
# constructor. Make every value that we also forward to PythonBackend
# an explicit argument so it is consumed exactly once instead of being
# duplicated inside **kwargs.
self.vocab_file = vocab_file
kwargs.setdefault("split_special_tokens",True)
super().__init__(
pad_token=pad_token,
bos_token=bos_token,
eos_token=eos_token,
unk_token=unk_token,
model_max_length=model_max_length,
padding_side=padding_side,
clean_up_tokenization_spaces=clean_up_tokenization_spaces,
**kwargs,
)
@property
def vocab_size(self):
return 259
def get_vocab(self):
vocab = {f"<0x{i:02X}>": i for i in range(256)}
vocab[PAD_TOKEN] = 256
vocab[BOS_TOKEN] = 257
vocab[EOS_TOKEN] = 258
return vocab
def _tokenize(self, text, **kwargs):
return [f"<0x{b:02X}>" for b in text.encode("utf-8", errors="replace")]
def _convert_token_to_id(self, token):
if token == PAD_TOKEN:
return 256
if token == BOS_TOKEN:
return 257
if token == EOS_TOKEN:
return 258
if isinstance(token, str) and token.startswith(BYTE_PREFIX) and token.endswith(">"):
try:
value = int(token[3:-1], 16)
if 0 <= value <= 255:
return value
except ValueError:
pass
# This branch should be unreachable for text encoded by this tokenizer.
return 0
def _convert_id_to_token(self, index):
index = int(index)
if 0 <= index <= 255:
return f"<0x{index:02X}>"
if index == 256:
return PAD_TOKEN
if index == 257:
return BOS_TOKEN
if index == 258:
return EOS_TOKEN
return "<0x00>"
def convert_tokens_to_string(self, tokens):
out = []
buf = bytearray()
def flush():
nonlocal buf
if buf:
out.append(bytes(buf).decode("utf-8", errors="replace"))
buf = bytearray()
for token in tokens:
idx = self._convert_token_to_id(token)
if isinstance(token, str) and 0 <= idx <= 255 and token.startswith(BYTE_PREFIX):
buf.append(idx)
else:
flush()
out.append(str(token))
flush()
return "".join(out)
def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
# BET pretraining did not automatically insert BOS/EOS around ordinary text.
if token_ids_1 is None:
return list(token_ids_0)
return list(token_ids_0) + list(token_ids_1)
def create_token_type_ids_from_sequences(self, token_ids_0, token_ids_1=None):
n = len(token_ids_0) + (len(token_ids_1) if token_ids_1 is not None else 0)
return [0] * n
def save_vocabulary(self, save_directory, filename_prefix=None):
os.makedirs(save_directory, exist_ok=True)
name = "byte_vocab.json" if filename_prefix is None else f"{filename_prefix}-byte_vocab.json"
path = os.path.join(save_directory, name)
vocab = {f"<0x{i:02X}>": i for i in range(256)}
vocab.update({PAD_TOKEN: 256, BOS_TOKEN: 257, EOS_TOKEN: 258})
with open(path, "w", encoding="utf-8") as f:
json.dump(vocab, f, indent=2, sort_keys=True)
return (path,)
|