Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 20243#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""Tokenization class for model MyT5."""16 17import json18import os19import warnings20from collections import defaultdict21from typing import Optional, Union22 23from ...tokenization_utils import AddedToken, PreTrainedTokenizer24from ...utils import logging25 26 27logger = logging.get_logger(__name__)28 29 30VOCAB_FILES_NAMES = {"vocab_file": "byte_maps.json"}31 32 33class ByteRewriter:34 """35 Byte rewriter class for MyT5 tokenizer.36 This class is used to rewrite bytes using a hash tree. The hash tree is constructed from a set of rewriting rules.37 38 Args:39 rewriting_rules (`str` or `dict[str, str]`):40 A path to a json file containing the rewriting rules or a dictionary containing the rewriting rules.41 42 """43 44 LEAF = "[LEAF]"45 46 def __init__(self, rewriting_rules: Union[str, dict[str, str]]):47 if isinstance(rewriting_rules, str):48 with open(rewriting_rules, "r") as f:49 rewriting_rules = json.load(f)50 elif not isinstance(rewriting_rules, dict):51 raise TypeError(52 f"rewriting_rules should be either a path to json file or a dict, got {type(rewriting_rules)}"53 )54 55 self.hash_tree = self.construct_hash_tree(rewriting_rules)56 reverse_rewriting_rules = {v: k for k, v in rewriting_rules.items()}57 self.reverse_hash_tree = self.construct_hash_tree(reverse_rewriting_rules)58 59 def add_leaf(self, hash_tree: dict[str, Union[dict, list[str]]], byte_in_sequence: str, byte_out_sequence: str):60 """61 Add a leaf with the output byte sequence to the hash tree.62 """63 byte_in_list = byte_in_sequence.split(" ")64 byte_out_list = byte_out_sequence.split(" ")65 66 tree_pointer = hash_tree67 for b in byte_in_list:68 if b not in tree_pointer:69 tree_pointer[b] = {}70 tree_pointer = tree_pointer[b]71 72 tree_pointer[self.LEAF] = byte_out_list73 74 def construct_hash_tree(self, rewriting_rules: dict[str, str]) -> dict[str, Union[dict, list[str]]]:75 """76 Construct a hash tree for rewritten byte sequences.77 """78 hash_tree = defaultdict(dict)79 for b in (f"{x:02x}" for x in range(256)):80 hash_tree[b][self.LEAF] = [b]81 82 for in_sequence, out_sequence in rewriting_rules.items():83 self.add_leaf(hash_tree, in_sequence, out_sequence)84 85 return hash_tree86 87 def search_hash_tree(self, byte_sequence: list[str]) -> Union[None, list[str]]:88 """89 Search the hash tree and return the rewritten byte sequence if found.90 """91 tree_pointer = self.hash_tree92 for b in byte_sequence:93 if b in tree_pointer:94 tree_pointer = tree_pointer[b]95 else:96 return None97 98 return tree_pointer[self.LEAF]99 100 def rewrite_bytes(self, in_bytes: list[str], reverse=False) -> list[str]:101 """102 Rewrite a sequence of bytes using the hash tree.103 104 Args:105 in_bytes (`list[str]`): A list of bytes to be rewritten.106 reverse (`bool`): If True, decoding is performed with the reverse hash tree.107 Returns:108 `list[str]`: The rewritten byte sequence.109 """110 out_bytes = []111 b_start = 0112 b_end = 0113 114 while b_start < len(in_bytes):115 tree_pointer = self.hash_tree if not reverse else self.reverse_hash_tree116 for j in range(b_start, len(in_bytes)):117 b = in_bytes[j]118 if b in tree_pointer:119 tree_pointer = tree_pointer[b]120 elif j == b_start:121 cur_leaf = [b]122 b_end = j123 break124 else:125 break126 if self.LEAF in tree_pointer:127 cur_leaf = tree_pointer[self.LEAF]128 b_end = j129 out_bytes.extend(cur_leaf)130 b_start = b_end + 1131 132 return out_bytes133 134 135class MyT5Tokenizer(PreTrainedTokenizer):136 """137 Construct a MyT5 tokenizer.138 139 This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to140 this superclass for more information regarding those methods.141 142 Args:143 vocab_file (`str`): The file containing the byte rewriting rules.144 eos_token (`str`, *optional*, defaults to `"</s>"`):145 The end of sequence token.146 147 unk_token (`str`, *optional*, defaults to `"<unk>"`):148 The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this149 token instead.150 pad_token (`str`, *optional*, defaults to `"<pad>"`):151 The token used for padding, for example when batching sequences of different lengths.152 extra_ids (`int`, *optional*, defaults to 125):153 Add a number of extra ids added to the end of the vocabulary for use as sentinels. These tokens are154 accessible as "<extra_id_{%d}>" where "{%d}" is a number between 0 and extra_ids-1. Extra tokens are155 indexed from the end of the vocabulary up to beginning ("<extra_id_0>" is the last token in the vocabulary156 like in ByT5 preprocessing see157 [here](https://github.com/google-research/text-to-text-transfer-transformer/blob/9fd7b14a769417be33bc6c850f9598764913c833/t5/data/preprocessors.py#L2117)).158 additional_special_tokens (`list[str]`, *optional*):159 Additional special tokens used by the tokenizer.160 """161 162 model_input_names = ["input_ids", "attention_mask"]163 vocab_files_names = VOCAB_FILES_NAMES164 165 def __init__(166 self,167 vocab_file,168 eos_token="</s>",169 unk_token="<unk>",170 pad_token="<pad>",171 extra_ids=125,172 additional_special_tokens=None,173 **kwargs,174 ) -> None:175 # Add extra_ids to the special token list176 if extra_ids > 0 and additional_special_tokens is None:177 additional_special_tokens = [f"<extra_id_{i}>" for i in range(extra_ids)]178 elif extra_ids > 0 and additional_special_tokens is not None and len(additional_special_tokens) > 0:179 # Check that we have the right number of extra_id special tokens180 extra_tokens = len(set(filter(lambda x: bool("extra_id" in str(x)), additional_special_tokens)))181 if extra_tokens != extra_ids:182 raise ValueError(183 f"Both extra_ids ({extra_ids}) and additional_special_tokens ({additional_special_tokens}) are"184 " provided to MyT5Tokenizer. In this case the additional_special_tokens must include the"185 " extra_ids tokens"186 )187 188 pad_token = AddedToken(pad_token, lstrip=True, rstrip=True) if isinstance(pad_token, str) else pad_token189 eos_token = AddedToken(eos_token, lstrip=True, rstrip=True) if isinstance(eos_token, str) else eos_token190 unk_token = AddedToken(unk_token, lstrip=True, rstrip=True) if isinstance(unk_token, str) else unk_token191 # unk token needs to be in the vocab with correct index192 self._added_tokens_decoder = {0: pad_token, 1: eos_token, 2: unk_token}193 self.offset = len(self._added_tokens_decoder)194 self._utf_vocab_size = 2**8 # utf is 8 bits195 196 # Load byte maps197 self.byte_maps = json.load(open(vocab_file, "r"))198 199 self.decompose_rewriter = ByteRewriter(self.byte_maps["decompose_map"])200 self.merge_rewriter = ByteRewriter(self.byte_maps["merge_map"])201 202 super().__init__(203 eos_token=eos_token,204 unk_token=unk_token,205 pad_token=pad_token,206 extra_ids=0,207 additional_special_tokens=additional_special_tokens,208 **kwargs,209 )210 211 @property212 def vocab_size(self):213 return self._utf_vocab_size214 215 # Copied from transformers.models.byt5.tokenization_byt5.ByT5Tokenizer.get_vocab216 def get_vocab(self):217 vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size + self.offset)}218 vocab.update(self.added_tokens_encoder)219 return vocab220 221 # Copied from transformers.models.byt5.tokenization_byt5.ByT5Tokenizer.get_special_tokens_mask222 def get_special_tokens_mask(223 self, token_ids_0: list[int], token_ids_1: Optional[list[int]] = None, already_has_special_tokens: bool = False224 ) -> list[int]:225 """226 Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding227 special tokens using the tokenizer `prepare_for_model` method.228 229 Args:230 token_ids_0 (`list[int]`):231 List of IDs.232 token_ids_1 (`list[int]`, *optional*):233 Optional second list of IDs for sequence pairs.234 already_has_special_tokens (`bool`, *optional*, defaults to `False`):235 Whether or not the token list is already formatted with special tokens for the model.236 237 Returns:238 `list[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.239 """240 if already_has_special_tokens:241 return super().get_special_tokens_mask(242 token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True243 )244 245 # normal case: some special tokens246 if token_ids_1 is None:247 return ([0] * len(token_ids_0)) + [1]248 return ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]249 250 def _add_eos_if_not_present(self, token_ids: list[int]) -> list[int]:251 """Do not add eos again if user already added it."""252 if len(token_ids) > 0 and token_ids[-1] == self.eos_token_id:253 warnings.warn(254 f"This sequence already has {self.eos_token}. In future versions this behavior may lead to duplicated"255 " eos tokens being added."256 )257 return token_ids258 else:259 return token_ids + [self.eos_token_id]260 261 def create_token_type_ids_from_sequences(262 self, token_ids_0: list[int], token_ids_1: Optional[list[int]] = None263 ) -> list[int]:264 """265 Create a mask from the two sequences passed to be used in a sequence-pair classification task. MyT5 does not266 make use of token type ids, therefore a list of zeros is returned.267 268 Args:269 token_ids_0 (`list[int]`):270 List of IDs.271 token_ids_1 (`list[int]`, *optional*):272 Optional second list of IDs for sequence pairs.273 274 Returns:275 `list[int]`: List of zeros.276 """277 eos = [self.eos_token_id]278 279 if token_ids_1 is None:280 return len(token_ids_0 + eos) * [0]281 return len(token_ids_0 + eos + token_ids_1 + eos) * [0]282 283 # Copied from transformers.models.byt5.tokenization_byt5.ByT5Tokenizer.build_inputs_with_special_tokens284 def build_inputs_with_special_tokens(285 self, token_ids_0: list[int], token_ids_1: Optional[list[int]] = None286 ) -> list[int]:287 """288 Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and289 adding special tokens. A sequence has the following format:290 291 - single sequence: `X </s>`292 - pair of sequences: `A </s> B </s>`293 294 Args:295 token_ids_0 (`list[int]`):296 List of IDs to which the special tokens will be added.297 token_ids_1 (`list[int]`, *optional*):298 Optional second list of IDs for sequence pairs.299 300 Returns:301 `list[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.302 """303 token_ids_0 = self._add_eos_if_not_present(token_ids_0)304 if token_ids_1 is None:305 return token_ids_0306 else:307 token_ids_1 = self._add_eos_if_not_present(token_ids_1)308 return token_ids_0 + token_ids_1309 310 def _tokenize(self, text: str, **kwargs) -> list[str]:311 """Take as input a string and return a list of strings (tokens) for words/sub-words.312 Represents tokens in two character hex format"""313 314 tokens = [f"{i:02x}" for i in text.encode("utf-8")]315 tokens = self.morphological_encode(tokens)316 return tokens317 318 def _convert_token_to_id(self, token):319 """Converts a token (str) in an id using the vocab."""320 321 if len(token) != 2:322 token_id = None323 else:324 token_id = int(token, 16) + self.offset325 326 return token_id327 328 def _convert_id_to_token(self, index):329 """Converts an index (integer) in a token (str) using the vocab."""330 token = f"{index - self.offset:02x}"331 return token332 333 def morphological_encode(self, indices: list[str]) -> list[str]:334 # Decompose and merge morphological sequences335 indices = self.decompose_rewriter.rewrite_bytes(indices, reverse=False)336 indices = self.merge_rewriter.rewrite_bytes(indices, reverse=False)337 return indices338 339 def morphological_decode(self, indices: list[str]) -> list[str]:340 # Demerge and compose morphological sequences341 indices = self.merge_rewriter.rewrite_bytes(indices, reverse=True)342 indices = self.decompose_rewriter.rewrite_bytes(indices, reverse=True)343 return indices344 345 def convert_tokens_to_string(self, tokens):346 """Converts a sequence of tokens (string) in a single string."""347 bstring = b""348 349 out_tokens = []350 for token in tokens:351 if token in self.added_tokens_decoder:352 out_tokens.append(self.added_tokens_decoder[token])353 elif token in self.added_tokens_encoder:354 out_tokens.append(token)355 else:356 out_tokens.append(token)357 358 out_tokens = self.morphological_decode(out_tokens)359 _added_tokens = set(self.added_tokens_decoder.values()) | set(self.added_tokens_encoder)360 for token in out_tokens:361 if token in _added_tokens:362 bstring += bytes(token, "utf-8")363 else:364 bstring += bytes.fromhex(token)365 string = bstring.decode("utf-8", errors="ignore")366 return string367 368 def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> tuple[str]:369 if os.path.isdir(save_directory):370 vocab_file = os.path.join(371 save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]372 )373 else:374 vocab_file = (filename_prefix + "-" if filename_prefix else "") + save_directory375 with open(vocab_file, "w", encoding="utf-8") as writer:376 writer.write(json.dumps(self.byte_maps, indent=2, ensure_ascii=False))377 return (vocab_file,)378 379 380__all__ = ["MyT5Tokenizer"]381