File size: 2,898 Bytes
29a4421 | 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 139 140 141 |
import json
from pathlib import Path
from transformers import PreTrainedTokenizer
class SupernovaTokenizer(PreTrainedTokenizer):
vocab_files_names = {
"vocab_file": "vocab.json"
}
model_input_names = [
"input_ids",
"attention_mask",
]
def __init__(
self,
vocab_file,
bos_token="<BOS>",
eos_token="<EOS>",
pad_token="<PAD>",
user_token="<USER>",
assistant_token="<ASSISTANT>",
**kwargs,
):
with open(
vocab_file,
"r",
encoding="utf-8",
) as f:
vocab = json.load(f)
self.vocab = {
str(token): int(idx)
for token, idx in vocab.items()
}
self.ids_to_tokens = {
int(idx): str(token)
for token, idx in self.vocab.items()
}
self.user_token = user_token
self.assistant_token = assistant_token
super().__init__(
bos_token=bos_token,
eos_token=eos_token,
pad_token=pad_token,
**kwargs,
)
@property
def vocab_size(self):
return len(self.vocab)
def get_vocab(self):
return dict(self.vocab)
def _tokenize(self, text, **kwargs):
return list(text)
def _convert_token_to_id(self, token):
if token not in self.vocab:
raise ValueError(
f"Unknown character/token: {token!r}"
)
return self.vocab[token]
def _convert_id_to_token(self, index):
index = int(index)
if index not in self.ids_to_tokens:
raise ValueError(
f"Unknown token id: {index}"
)
return self.ids_to_tokens[index]
def convert_tokens_to_string(self, tokens):
return "".join(tokens)
def build_inputs_with_special_tokens(
self,
token_ids_0,
token_ids_1=None,
):
result = [self.bos_token_id]
result += list(token_ids_0)
if token_ids_1 is not None:
result += list(token_ids_1)
result += [self.eos_token_id]
return result
def save_vocabulary(
self,
save_directory,
filename_prefix=None,
):
save_directory = Path(save_directory)
save_directory.mkdir(
parents=True,
exist_ok=True,
)
filename = "vocab.json"
if filename_prefix:
filename = (
f"{filename_prefix}-vocab.json"
)
path = save_directory / filename
with open(
path,
"w",
encoding="utf-8",
) as f:
json.dump(
self.vocab,
f,
ensure_ascii=False,
indent=2,
)
return (str(path),)
|