johnpaulbin commited on
Commit
5b1cec5
1 Parent(s): 4210065

Upload tokenization_small100.py

Browse files
Files changed (1) hide show
  1. tokenization_small100.py +365 -0
tokenization_small100.py ADDED
@@ -0,0 +1,365 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2022 Idiap Research Institute, http://www.idiap.ch/
2
+ # Written by Alireza Mohammadshahi <alireza.mohammadshahi@idiap.ch>
3
+ # This is a modified version of https://github.com/huggingface/transformers/blob/main/src/transformers/models/m2m_100/tokenization_m2m_100.py
4
+ # which owns by Fariseq Authors and The HuggingFace Inc. team.
5
+ #
6
+ #
7
+ # Licensed under the Apache License, Version 2.0 (the "License");
8
+ # you may not use this file except in compliance with the License.
9
+ # You may obtain a copy of the License at
10
+ #
11
+ # http://www.apache.org/licenses/LICENSE-2.0
12
+ #
13
+ # Unless required by applicable law or agreed to in writing, software
14
+ # distributed under the License is distributed on an "AS IS" BASIS,
15
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16
+ # See the License for the specific language governing permissions and
17
+ # limitations under the License.
18
+ """Tokenization classes for SMALL100."""
19
+ import json
20
+ import os
21
+ from pathlib import Path
22
+ from shutil import copyfile
23
+ from typing import Any, Dict, List, Optional, Tuple, Union
24
+
25
+ import sentencepiece
26
+
27
+ from transformers.tokenization_utils import BatchEncoding, PreTrainedTokenizer
28
+ from transformers.utils import logging
29
+
30
+
31
+ logger = logging.get_logger(__name__)
32
+
33
+ SPIECE_UNDERLINE = "▁"
34
+
35
+ VOCAB_FILES_NAMES = {
36
+ "vocab_file": "vocab.json",
37
+ "spm_file": "sentencepiece.bpe.model",
38
+ "tokenizer_config_file": "tokenizer_config.json",
39
+ }
40
+
41
+ PRETRAINED_VOCAB_FILES_MAP = {
42
+ "vocab_file": {
43
+ "alirezamsh/small100": "https://huggingface.co/alirezamsh/small100/resolve/main/vocab.json",
44
+ },
45
+ "spm_file": {
46
+ "alirezamsh/small100": "https://huggingface.co/alirezamsh/small100/resolve/main/sentencepiece.bpe.model",
47
+ },
48
+ "tokenizer_config_file": {
49
+ "alirezamsh/small100": "https://huggingface.co/alirezamsh/small100/resolve/main/tokenizer_config.json",
50
+ },
51
+ }
52
+
53
+ PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
54
+ "alirezamsh/small100": 1024,
55
+ }
56
+
57
+ # fmt: off
58
+ FAIRSEQ_LANGUAGE_CODES = {
59
+ "m2m100": ["af", "am", "ar", "ast", "az", "ba", "be", "bg", "bn", "br", "bs", "ca", "ceb", "cs", "cy", "da", "de", "el", "en", "es", "et", "fa", "ff", "fi", "fr", "fy", "ga", "gd", "gl", "gu", "ha", "he", "hi", "hr", "ht", "hu", "hy", "id", "ig", "ilo", "is", "it", "ja", "jv", "ka", "kk", "km", "kn", "ko", "lb", "lg", "ln", "lo", "lt", "lv", "mg", "mk", "ml", "mn", "mr", "ms", "my", "ne", "nl", "no", "ns", "oc", "or", "pa", "pl", "ps", "pt", "ro", "ru", "sd", "si", "sk", "sl", "so", "sq", "sr", "ss", "su", "sv", "sw", "ta", "th", "tl", "tn", "tr", "uk", "ur", "uz", "vi", "wo", "xh", "yi", "yo", "zh", "zu"]
60
+ }
61
+ # fmt: on
62
+
63
+
64
+ class SMALL100Tokenizer(PreTrainedTokenizer):
65
+ """
66
+ Construct an SMALL100 tokenizer. Based on [SentencePiece](https://github.com/google/sentencepiece).
67
+ This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
68
+ this superclass for more information regarding those methods.
69
+ Args:
70
+ vocab_file (`str`):
71
+ Path to the vocabulary file.
72
+ spm_file (`str`):
73
+ Path to [SentencePiece](https://github.com/google/sentencepiece) file (generally has a .spm extension) that
74
+ contains the vocabulary.
75
+ tgt_lang (`str`, *optional*):
76
+ A string representing the target language.
77
+ eos_token (`str`, *optional*, defaults to `"</s>"`):
78
+ The end of sequence token.
79
+ sep_token (`str`, *optional*, defaults to `"</s>"`):
80
+ The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
81
+ sequence classification or for a text and a question for question answering. It is also used as the last
82
+ token of a sequence built with special tokens.
83
+ unk_token (`str`, *optional*, defaults to `"<unk>"`):
84
+ The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
85
+ token instead.
86
+ pad_token (`str`, *optional*, defaults to `"<pad>"`):
87
+ The token used for padding, for example when batching sequences of different lengths.
88
+ language_codes (`str`, *optional*):
89
+ What language codes to use. Should be `"m2m100"`.
90
+ sp_model_kwargs (`dict`, *optional*):
91
+ Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for
92
+ SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things,
93
+ to set:
94
+ - `enable_sampling`: Enable subword regularization.
95
+ - `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout.
96
+ - `nbest_size = {0,1}`: No sampling is performed.
97
+ - `nbest_size > 1`: samples from the nbest_size results.
98
+ - `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice)
99
+ using forward-filtering-and-backward-sampling algorithm.
100
+ - `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for
101
+ BPE-dropout.
102
+ Examples:
103
+ ```python
104
+ >>> from tokenization_small100 import SMALL100Tokenizer
105
+ >>> tokenizer = SMALL100Tokenizer.from_pretrained("alirezamsh/small100", tgt_lang="ro")
106
+ >>> src_text = " UN Chief Says There Is No Military Solution in Syria"
107
+ >>> tgt_text = "Şeful ONU declară că nu există o soluţie militară în Siria"
108
+ >>> model_inputs = tokenizer(src_text, text_target=tgt_text, return_tensors="pt")
109
+ >>> model(**model_inputs) # should work
110
+ ```"""
111
+
112
+ vocab_files_names = VOCAB_FILES_NAMES
113
+ max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
114
+ pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
115
+ model_input_names = ["input_ids", "attention_mask"]
116
+
117
+ prefix_tokens: List[int] = []
118
+ suffix_tokens: List[int] = []
119
+
120
+ def __init__(
121
+ self,
122
+ vocab_file,
123
+ spm_file,
124
+ tgt_lang=None,
125
+ bos_token="<s>",
126
+ eos_token="</s>",
127
+ sep_token="</s>",
128
+ pad_token="<pad>",
129
+ unk_token="<unk>",
130
+ language_codes="m2m100",
131
+ sp_model_kwargs: Optional[Dict[str, Any]] = None,
132
+ num_madeup_words=8,
133
+ **kwargs,
134
+ ) -> None:
135
+ self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
136
+
137
+ self.language_codes = language_codes
138
+ fairseq_language_code = FAIRSEQ_LANGUAGE_CODES[language_codes]
139
+ self.lang_code_to_token = {lang_code: f"__{lang_code}__" for lang_code in fairseq_language_code}
140
+
141
+ kwargs["additional_special_tokens"] = kwargs.get("additional_special_tokens", [])
142
+ kwargs["additional_special_tokens"] += [
143
+ self.get_lang_token(lang_code)
144
+ for lang_code in fairseq_language_code
145
+ if self.get_lang_token(lang_code) not in kwargs["additional_special_tokens"]
146
+ ]
147
+
148
+ self.vocab_file = vocab_file
149
+ self.encoder = load_json(vocab_file)
150
+ self.decoder = {v: k for k, v in self.encoder.items()}
151
+ self.spm_file = spm_file
152
+ self.sp_model = load_spm(spm_file, self.sp_model_kwargs)
153
+
154
+ self.encoder_size = len(self.encoder)
155
+
156
+ self.lang_token_to_id = {
157
+ self.get_lang_token(lang_code): self.encoder_size + i for i, lang_code in enumerate(fairseq_language_code)
158
+ }
159
+ self.lang_code_to_id = {lang_code: self.encoder_size + i for i, lang_code in enumerate(fairseq_language_code)}
160
+ self.id_to_lang_token = {v: k for k, v in self.lang_token_to_id.items()}
161
+
162
+ self._tgt_lang = tgt_lang if tgt_lang is not None else "en"
163
+ self.cur_lang_id = self.get_lang_id(self._tgt_lang)
164
+ self.num_madeup_words = num_madeup_words
165
+
166
+ super().__init__(
167
+ tgt_lang=tgt_lang,
168
+ bos_token=bos_token,
169
+ eos_token=eos_token,
170
+ sep_token=sep_token,
171
+ unk_token=unk_token,
172
+ pad_token=pad_token,
173
+ language_codes=language_codes,
174
+ sp_model_kwargs=self.sp_model_kwargs,
175
+ num_madeup_words=num_madeup_words,
176
+ **kwargs,
177
+ )
178
+
179
+ self.set_lang_special_tokens(self._tgt_lang)
180
+
181
+
182
+ @property
183
+ def vocab_size(self) -> int:
184
+ return len(self.encoder) + len(self.lang_token_to_id) + self.num_madeup_words
185
+
186
+ @property
187
+ def tgt_lang(self) -> str:
188
+ return self._tgt_lang
189
+
190
+ @tgt_lang.setter
191
+ def tgt_lang(self, new_tgt_lang: str) -> None:
192
+ self._tgt_lang = new_tgt_lang
193
+ self.set_lang_special_tokens(self._tgt_lang)
194
+
195
+ def _tokenize(self, text: str) -> List[str]:
196
+ return self.sp_model.encode(text, out_type=str)
197
+
198
+ def _convert_token_to_id(self, token):
199
+ if token in self.lang_token_to_id:
200
+ return self.lang_token_to_id[token]
201
+ return self.encoder.get(token, self.encoder[self.unk_token])
202
+
203
+ def _convert_id_to_token(self, index: int) -> str:
204
+ """Converts an index (integer) in a token (str) using the decoder."""
205
+ if index in self.id_to_lang_token:
206
+ return self.id_to_lang_token[index]
207
+ return self.decoder.get(index, self.unk_token)
208
+
209
+ def convert_tokens_to_string(self, tokens: List[str]) -> str:
210
+ """Converts a sequence of tokens (strings for sub-words) in a single string."""
211
+ return self.sp_model.decode(tokens)
212
+
213
+ def get_special_tokens_mask(
214
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
215
+ ) -> List[int]:
216
+ """
217
+ Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
218
+ special tokens using the tokenizer `prepare_for_model` method.
219
+ Args:
220
+ token_ids_0 (`List[int]`):
221
+ List of IDs.
222
+ token_ids_1 (`List[int]`, *optional*):
223
+ Optional second list of IDs for sequence pairs.
224
+ already_has_special_tokens (`bool`, *optional*, defaults to `False`):
225
+ Whether or not the token list is already formatted with special tokens for the model.
226
+ Returns:
227
+ `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
228
+ """
229
+
230
+ if already_has_special_tokens:
231
+ return super().get_special_tokens_mask(
232
+ token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
233
+ )
234
+
235
+ prefix_ones = [1] * len(self.prefix_tokens)
236
+ suffix_ones = [1] * len(self.suffix_tokens)
237
+ if token_ids_1 is None:
238
+ return prefix_ones + ([0] * len(token_ids_0)) + suffix_ones
239
+ return prefix_ones + ([0] * len(token_ids_0)) + ([0] * len(token_ids_1)) + suffix_ones
240
+
241
+ def build_inputs_with_special_tokens(
242
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
243
+ ) -> List[int]:
244
+ """
245
+ Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
246
+ adding special tokens. An MBART sequence has the following format, where `X` represents the sequence:
247
+ - `input_ids` (for encoder) `X [eos, src_lang_code]`
248
+ - `decoder_input_ids`: (for decoder) `X [eos, tgt_lang_code]`
249
+ BOS is never used. Pairs of sequences are not the expected use case, but they will be handled without a
250
+ separator.
251
+ Args:
252
+ token_ids_0 (`List[int]`):
253
+ List of IDs to which the special tokens will be added.
254
+ token_ids_1 (`List[int]`, *optional*):
255
+ Optional second list of IDs for sequence pairs.
256
+ Returns:
257
+ `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
258
+ """
259
+ if token_ids_1 is None:
260
+ if self.prefix_tokens is None:
261
+ return token_ids_0 + self.suffix_tokens
262
+ else:
263
+ return self.prefix_tokens + token_ids_0 + self.suffix_tokens
264
+ # We don't expect to process pairs, but leave the pair logic for API consistency
265
+ if self.prefix_tokens is None:
266
+ return token_ids_0 + token_ids_1 + self.suffix_tokens
267
+ else:
268
+ return self.prefix_tokens + token_ids_0 + token_ids_1 + self.suffix_tokens
269
+
270
+ def get_vocab(self) -> Dict:
271
+ vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
272
+ vocab.update(self.added_tokens_encoder)
273
+ return vocab
274
+
275
+ def __getstate__(self) -> Dict:
276
+ state = self.__dict__.copy()
277
+ state["sp_model"] = None
278
+ return state
279
+
280
+ def __setstate__(self, d: Dict) -> None:
281
+ self.__dict__ = d
282
+
283
+ # for backward compatibility
284
+ if not hasattr(self, "sp_model_kwargs"):
285
+ self.sp_model_kwargs = {}
286
+
287
+ self.sp_model = load_spm(self.spm_file, self.sp_model_kwargs)
288
+
289
+ def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
290
+ save_dir = Path(save_directory)
291
+ if not save_dir.is_dir():
292
+ raise OSError(f"{save_directory} should be a directory")
293
+ vocab_save_path = save_dir / (
294
+ (filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["vocab_file"]
295
+ )
296
+ spm_save_path = save_dir / (
297
+ (filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["spm_file"]
298
+ )
299
+
300
+ save_json(self.encoder, vocab_save_path)
301
+
302
+ if os.path.abspath(self.spm_file) != os.path.abspath(spm_save_path) and os.path.isfile(self.spm_file):
303
+ copyfile(self.spm_file, spm_save_path)
304
+ elif not os.path.isfile(self.spm_file):
305
+ with open(spm_save_path, "wb") as fi:
306
+ content_spiece_model = self.sp_model.serialized_model_proto()
307
+ fi.write(content_spiece_model)
308
+
309
+ return (str(vocab_save_path), str(spm_save_path))
310
+
311
+ def prepare_seq2seq_batch(
312
+ self,
313
+ src_texts: List[str],
314
+ tgt_texts: Optional[List[str]] = None,
315
+ tgt_lang: str = "ro",
316
+ **kwargs,
317
+ ) -> BatchEncoding:
318
+ self.tgt_lang = tgt_lang
319
+ self.set_lang_special_tokens(self.tgt_lang)
320
+ return super().prepare_seq2seq_batch(src_texts, tgt_texts, **kwargs)
321
+
322
+ def _build_translation_inputs(self, raw_inputs, tgt_lang: Optional[str], **extra_kwargs):
323
+ """Used by translation pipeline, to prepare inputs for the generate function"""
324
+ if tgt_lang is None:
325
+ raise ValueError("Translation requires a `tgt_lang` for this model")
326
+ self.tgt_lang = tgt_lang
327
+ inputs = self(raw_inputs, add_special_tokens=True, **extra_kwargs)
328
+ return inputs
329
+
330
+ def _switch_to_input_mode(self):
331
+ self.set_lang_special_tokens(self.tgt_lang)
332
+
333
+ def _switch_to_target_mode(self):
334
+ self.prefix_tokens = None
335
+ self.suffix_tokens = [self.eos_token_id]
336
+
337
+ def set_lang_special_tokens(self, src_lang: str) -> None:
338
+ """Reset the special tokens to the tgt lang setting. No prefix and suffix=[eos, tgt_lang_code]."""
339
+ lang_token = self.get_lang_token(src_lang)
340
+ self.cur_lang_id = self.lang_token_to_id[lang_token]
341
+ self.prefix_tokens = [self.cur_lang_id]
342
+ self.suffix_tokens = [self.eos_token_id]
343
+
344
+ def get_lang_token(self, lang: str) -> str:
345
+ return self.lang_code_to_token[lang]
346
+
347
+ def get_lang_id(self, lang: str) -> int:
348
+ lang_token = self.get_lang_token(lang)
349
+ return self.lang_token_to_id[lang_token]
350
+
351
+
352
+ def load_spm(path: str, sp_model_kwargs: Dict[str, Any]) -> sentencepiece.SentencePieceProcessor:
353
+ spm = sentencepiece.SentencePieceProcessor(**sp_model_kwargs)
354
+ spm.Load(str(path))
355
+ return spm
356
+
357
+
358
+ def load_json(path: str) -> Union[Dict, List]:
359
+ with open(path, "r") as f:
360
+ return json.load(f)
361
+
362
+
363
+ def save_json(data, path: str) -> None:
364
+ with open(path, "w") as f:
365
+ json.dump(data, f, indent=2)