aapot commited on
Commit
adebfcb
1 Parent(s): bbbd016

Add tokenizer

Browse files
.gitattributes CHANGED
File without changes
config.json ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "RobertaForMaskedLM"
4
+ ],
5
+ "attention_probs_dropout_prob": 0.1,
6
+ "bos_token_id": 0,
7
+ "classifier_dropout": null,
8
+ "eos_token_id": 2,
9
+ "gradient_checkpointing": false,
10
+ "hidden_act": "gelu",
11
+ "hidden_dropout_prob": 0.1,
12
+ "hidden_size": 1024,
13
+ "initializer_range": 0.02,
14
+ "intermediate_size": 4096,
15
+ "layer_norm_eps": 1e-05,
16
+ "max_position_embeddings": 514,
17
+ "model_type": "roberta",
18
+ "num_attention_heads": 16,
19
+ "num_hidden_layers": 24,
20
+ "pad_token_id": 1,
21
+ "position_embedding_type": "absolute",
22
+ "transformers_version": "4.10.0.dev0",
23
+ "type_vocab_size": 1,
24
+ "use_cache": true,
25
+ "vocab_size": 50265
26
+ }
merges.txt ADDED
The diff for this file is too large to render. See raw diff
 
run_mlm_flax.py ADDED
@@ -0,0 +1,734 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python
2
+ # coding=utf-8
3
+ # Copyright 2021 The HuggingFace Team All rights reserved.
4
+ #
5
+ # Licensed under the Apache License, Version 2.0 (the "License");
6
+ # you may not use this file except in compliance with the License.
7
+ # You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing, software
12
+ # distributed under the License is distributed on an "AS IS" BASIS,
13
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14
+ # See the License for the specific language governing permissions and
15
+ # limitations under the License.
16
+ """
17
+ Fine-tuning the library models for masked language modeling (BERT, ALBERT, RoBERTa...) with whole word masking on a
18
+ text file or a dataset.
19
+
20
+ Here is the full list of checkpoints on the hub that can be fine-tuned by this script:
21
+ https://huggingface.co/models?filter=masked-lm
22
+ """
23
+ import logging
24
+ import os
25
+ import sys
26
+ import time
27
+ from dataclasses import dataclass, field
28
+
29
+ # You can also adapt this script on your own masked language modeling task. Pointers for this are left as comments.
30
+ from pathlib import Path
31
+ from typing import Dict, List, Optional, Tuple
32
+
33
+ import numpy as np
34
+ from datasets import load_dataset
35
+ from tqdm import tqdm
36
+
37
+ import flax
38
+ import jax
39
+ import jax.numpy as jnp
40
+ import optax
41
+ from flax import jax_utils, traverse_util
42
+ from flax.training import train_state
43
+ from flax.training.common_utils import get_metrics, onehot, shard
44
+ from transformers import (
45
+ CONFIG_MAPPING,
46
+ FLAX_MODEL_FOR_MASKED_LM_MAPPING,
47
+ AutoConfig,
48
+ AutoTokenizer,
49
+ FlaxAutoModelForMaskedLM,
50
+ HfArgumentParser,
51
+ PreTrainedTokenizerBase,
52
+ TensorType,
53
+ TrainingArguments,
54
+ is_tensorboard_available,
55
+ set_seed,
56
+ )
57
+
58
+
59
+ MODEL_CONFIG_CLASSES = list(FLAX_MODEL_FOR_MASKED_LM_MAPPING.keys())
60
+ MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
61
+
62
+
63
+ @dataclass
64
+ class ModelArguments:
65
+ """
66
+ Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
67
+ """
68
+
69
+ model_name_or_path: Optional[str] = field(
70
+ default=None,
71
+ metadata={
72
+ "help": "The model checkpoint for weights initialization."
73
+ "Don't set if you want to train a model from scratch."
74
+ },
75
+ )
76
+ model_type: Optional[str] = field(
77
+ default=None,
78
+ metadata={"help": "If training from scratch, pass a model type from the list: " + ", ".join(MODEL_TYPES)},
79
+ )
80
+ config_name: Optional[str] = field(
81
+ default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
82
+ )
83
+ tokenizer_name: Optional[str] = field(
84
+ default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
85
+ )
86
+ cache_dir: Optional[str] = field(
87
+ default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
88
+ )
89
+ use_fast_tokenizer: bool = field(
90
+ default=True,
91
+ metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
92
+ )
93
+ dtype: Optional[str] = field(
94
+ default="float32",
95
+ metadata={
96
+ "help": "Floating-point format in which the model weights should be initialized and trained. Choose one of `[float32, float16, bfloat16]`."
97
+ },
98
+ )
99
+
100
+
101
+ @dataclass
102
+ class DataTrainingArguments:
103
+ """
104
+ Arguments pertaining to what data we are going to input our model for training and eval.
105
+ """
106
+
107
+ dataset_name: Optional[str] = field(
108
+ default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
109
+ )
110
+ dataset_config_name: Optional[str] = field(
111
+ default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
112
+ )
113
+ train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
114
+ dataset_filepath: Optional[str] = field(
115
+ default=None, metadata={"help": "Filepath to locally saved HF Dataset (with 'dataset.save_to_disk' method) to use for training"}
116
+ )
117
+ save_tokenized_dataset_filepath: Optional[str] = field(
118
+ default=None, metadata={"help": "Filepath for saving tokenized HF Dataset (with 'dataset.save_to_disk' method) to use for future trainings"}
119
+ )
120
+ tokenized_dataset_filepath: Optional[str] = field(
121
+ default=None, metadata={"help": "Filepath to locally saved tokenized HF Dataset (with 'dataset.save_to_disk' method) to use for training"}
122
+ )
123
+ validation_file: Optional[str] = field(
124
+ default=None,
125
+ metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},
126
+ )
127
+ train_ref_file: Optional[str] = field(
128
+ default=None,
129
+ metadata={"help": "An optional input train ref data file for whole word masking in Chinese."},
130
+ )
131
+ validation_ref_file: Optional[str] = field(
132
+ default=None,
133
+ metadata={"help": "An optional input validation ref data file for whole word masking in Chinese."},
134
+ )
135
+ overwrite_cache: bool = field(
136
+ default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
137
+ )
138
+ validation_split_percentage: Optional[int] = field(
139
+ default=5,
140
+ metadata={
141
+ "help": "The percentage of the train set used as validation set in case there's no validation split"
142
+ },
143
+ )
144
+ max_seq_length: Optional[int] = field(
145
+ default=None,
146
+ metadata={
147
+ "help": "The maximum total input sequence length after tokenization. Sequences longer "
148
+ "than this will be truncated. Default to the max input length of the model."
149
+ },
150
+ )
151
+ preprocessing_num_workers: Optional[int] = field(
152
+ default=None,
153
+ metadata={"help": "The number of processes to use for the preprocessing."},
154
+ )
155
+ mlm_probability: float = field(
156
+ default=0.15, metadata={"help": "Ratio of tokens to mask for masked language modeling loss"}
157
+ )
158
+ pad_to_max_length: bool = field(
159
+ default=False,
160
+ metadata={
161
+ "help": "Whether to pad all samples to `max_seq_length`. "
162
+ "If False, will pad the samples dynamically when batching to the maximum length in the batch."
163
+ },
164
+ )
165
+ line_by_line: bool = field(
166
+ default=False,
167
+ metadata={"help": "Whether distinct lines of text in the dataset are to be handled as distinct sequences."},
168
+ )
169
+
170
+ def __post_init__(self):
171
+ if self.dataset_name is None and self.train_file is None and self.dataset_filepath is None and self.validation_file is None:
172
+ raise ValueError("Need either a dataset name or a training/validation file.")
173
+ else:
174
+ if self.train_file is not None:
175
+ extension = self.train_file.split(".")[-1]
176
+ assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, a json or a txt file."
177
+ if self.validation_file is not None:
178
+ extension = self.validation_file.split(".")[-1]
179
+ assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, a json or a txt file."
180
+
181
+
182
+ @flax.struct.dataclass
183
+ class FlaxDataCollatorForLanguageModeling:
184
+ """
185
+ Data collator used for language modeling. Inputs are dynamically padded to the maximum length of a batch if they
186
+ are not all of the same length.
187
+
188
+ Args:
189
+ tokenizer (:class:`~transformers.PreTrainedTokenizer` or :class:`~transformers.PreTrainedTokenizerFast`):
190
+ The tokenizer used for encoding the data.
191
+ mlm_probability (:obj:`float`, `optional`, defaults to 0.15):
192
+ The probability with which to (randomly) mask tokens in the input.
193
+
194
+ .. note::
195
+
196
+ For best performance, this data collator should be used with a dataset having items that are dictionaries or
197
+ BatchEncoding, with the :obj:`"special_tokens_mask"` key, as returned by a
198
+ :class:`~transformers.PreTrainedTokenizer` or a :class:`~transformers.PreTrainedTokenizerFast` with the
199
+ argument :obj:`return_special_tokens_mask=True`.
200
+ """
201
+
202
+ tokenizer: PreTrainedTokenizerBase
203
+ mlm_probability: float = 0.15
204
+
205
+ def __post_init__(self):
206
+ if self.tokenizer.mask_token is None:
207
+ raise ValueError(
208
+ "This tokenizer does not have a mask token which is necessary for masked language modeling. "
209
+ "You should pass `mlm=False` to train on causal language modeling instead."
210
+ )
211
+
212
+ def __call__(self, examples: List[Dict[str, np.ndarray]], pad_to_multiple_of: int) -> Dict[str, np.ndarray]:
213
+ # Handle dict or lists with proper padding and conversion to tensor.
214
+ batch = self.tokenizer.pad(examples, pad_to_multiple_of=pad_to_multiple_of, return_tensors=TensorType.NUMPY)
215
+
216
+ # If special token mask has been preprocessed, pop it from the dict.
217
+ special_tokens_mask = batch.pop("special_tokens_mask", None)
218
+
219
+ batch["input_ids"], batch["labels"] = self.mask_tokens(
220
+ batch["input_ids"], special_tokens_mask=special_tokens_mask
221
+ )
222
+ return batch
223
+
224
+ def mask_tokens(
225
+ self, inputs: np.ndarray, special_tokens_mask: Optional[np.ndarray]
226
+ ) -> Tuple[np.ndarray, np.ndarray]:
227
+ """
228
+ Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.
229
+ """
230
+ labels = inputs.copy()
231
+ # We sample a few tokens in each sequence for MLM training (with probability `self.mlm_probability`)
232
+ probability_matrix = np.full(labels.shape, self.mlm_probability)
233
+ special_tokens_mask = special_tokens_mask.astype("bool")
234
+
235
+ probability_matrix[special_tokens_mask] = 0.0
236
+ masked_indices = np.random.binomial(1, probability_matrix).astype("bool")
237
+ labels[~masked_indices] = -100 # We only compute loss on masked tokens
238
+
239
+ # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
240
+ indices_replaced = np.random.binomial(1, np.full(labels.shape, 0.8)).astype("bool") & masked_indices
241
+ inputs[indices_replaced] = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token)
242
+
243
+ # 10% of the time, we replace masked input tokens with random word
244
+ indices_random = np.random.binomial(1, np.full(labels.shape, 0.5)).astype("bool")
245
+ indices_random &= masked_indices & ~indices_replaced
246
+
247
+ random_words = np.random.randint(self.tokenizer.vocab_size, size=labels.shape, dtype="i4")
248
+ inputs[indices_random] = random_words[indices_random]
249
+
250
+ # The rest of the time (10% of the time) we keep the masked input tokens unchanged
251
+ return inputs, labels
252
+
253
+
254
+ def generate_batch_splits(samples_idx: jnp.ndarray, batch_size: int) -> jnp.ndarray:
255
+ num_samples = len(samples_idx)
256
+ samples_to_remove = num_samples % batch_size
257
+
258
+ if samples_to_remove != 0:
259
+ samples_idx = samples_idx[:-samples_to_remove]
260
+ sections_split = num_samples // batch_size
261
+ batch_idx = np.split(samples_idx, sections_split)
262
+ return batch_idx
263
+
264
+
265
+ def write_train_metric(summary_writer, train_metrics, train_time, step):
266
+ summary_writer.scalar("train_time", train_time, step)
267
+
268
+ train_metrics = get_metrics(train_metrics)
269
+ for key, vals in train_metrics.items():
270
+ tag = f"train_{key}"
271
+ for i, val in enumerate(vals):
272
+ summary_writer.scalar(tag, val, step - len(vals) + i + 1)
273
+
274
+
275
+ def write_eval_metric(summary_writer, eval_metrics, step):
276
+ for metric_name, value in eval_metrics.items():
277
+ summary_writer.scalar(f"eval_{metric_name}", value, step)
278
+
279
+
280
+ if __name__ == "__main__":
281
+ # See all possible arguments in src/transformers/training_args.py
282
+ # or by passing the --help flag to this script.
283
+ # We now keep distinct sets of args, for a cleaner separation of concerns.
284
+
285
+ parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
286
+ if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
287
+ # If we pass only one argument to the script and it's the path to a json file,
288
+ # let's parse it to get our arguments.
289
+ model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
290
+ else:
291
+ model_args, data_args, training_args = parser.parse_args_into_dataclasses()
292
+
293
+ if (
294
+ os.path.exists(training_args.output_dir)
295
+ and os.listdir(training_args.output_dir)
296
+ and training_args.do_train
297
+ and not training_args.overwrite_output_dir
298
+ ):
299
+ raise ValueError(
300
+ f"Output directory ({training_args.output_dir}) already exists and is not empty."
301
+ "Use --overwrite_output_dir to overcome."
302
+ )
303
+
304
+ # Setup logging
305
+ logging.basicConfig(
306
+ format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
307
+ level="NOTSET",
308
+ datefmt="[%X]",
309
+ )
310
+
311
+ # Log on each process the small summary:
312
+ logger = logging.getLogger(__name__)
313
+
314
+ # Set the verbosity to info of the Transformers logger (on main process only):
315
+ logger.info(f"Training/evaluation parameters {training_args}")
316
+
317
+ # Set seed before initializing model.
318
+ set_seed(training_args.seed)
319
+
320
+ # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
321
+ # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
322
+ # (the dataset will be downloaded automatically from the datasets Hub).
323
+ #
324
+ # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
325
+ # 'text' is found. You can easily tweak this behavior (see below).
326
+ #
327
+ # In distributed training, the load_dataset function guarantees that only one local process can concurrently
328
+ # download the dataset.
329
+ if data_args.dataset_name is not None:
330
+ # Downloading and loading a dataset from the hub.
331
+ datasets = load_dataset(data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir)
332
+
333
+ if "validation" not in datasets.keys():
334
+ datasets["validation"] = load_dataset(
335
+ data_args.dataset_name,
336
+ data_args.dataset_config_name,
337
+ split=f"train[:{data_args.validation_split_percentage}%]",
338
+ cache_dir=model_args.cache_dir,
339
+ )
340
+ datasets["train"] = load_dataset(
341
+ data_args.dataset_name,
342
+ data_args.dataset_config_name,
343
+ split=f"train[{data_args.validation_split_percentage}%:]",
344
+ cache_dir=model_args.cache_dir,
345
+ )
346
+ elif data_args.dataset_filepath is not None:
347
+ # Loading a dataset from local file.
348
+ datasets = load_from_disk(data_args.dataset_filepath)
349
+ if "validation" not in datasets.keys():
350
+ datasets = datasets.train_test_split(test_size=data_args.validation_split_percentage/100)
351
+ datasets["validation"] = datasets["test"]
352
+ del datasets["test"]
353
+
354
+ else:
355
+ data_files = {}
356
+ if data_args.train_file is not None:
357
+ data_files["train"] = data_args.train_file
358
+ if data_args.validation_file is not None:
359
+ data_files["validation"] = data_args.validation_file
360
+ extension = data_args.train_file.split(".")[-1]
361
+ if extension == "txt":
362
+ extension = "text"
363
+ datasets = load_dataset(extension, data_files=data_files, cache_dir=model_args.cache_dir)
364
+
365
+ if "validation" not in datasets.keys():
366
+ datasets["validation"] = load_dataset(
367
+ extension,
368
+ data_files=data_files,
369
+ split=f"train[:{data_args.validation_split_percentage}%]",
370
+ cache_dir=model_args.cache_dir,
371
+ )
372
+ datasets["train"] = load_dataset(
373
+ extension,
374
+ data_files=data_files,
375
+ split=f"train[{data_args.validation_split_percentage}%:]",
376
+ cache_dir=model_args.cache_dir,
377
+ )
378
+ # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
379
+ # https://huggingface.co/docs/datasets/loading_datasets.html.
380
+
381
+ # Load pretrained model and tokenizer
382
+
383
+ # Distributed training:
384
+ # The .from_pretrained methods guarantee that only one local process can concurrently
385
+ # download model & vocab.
386
+ if model_args.config_name:
387
+ config = AutoConfig.from_pretrained(model_args.config_name, cache_dir=model_args.cache_dir)
388
+ elif model_args.model_name_or_path:
389
+ config = AutoConfig.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir)
390
+ else:
391
+ config = CONFIG_MAPPING[model_args.model_type]()
392
+ logger.warning("You are instantiating a new config instance from scratch.")
393
+
394
+ if model_args.tokenizer_name:
395
+ tokenizer = AutoTokenizer.from_pretrained(
396
+ model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
397
+ )
398
+ elif model_args.model_name_or_path:
399
+ tokenizer = AutoTokenizer.from_pretrained(
400
+ model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
401
+ )
402
+ else:
403
+ raise ValueError(
404
+ "You are instantiating a new tokenizer from scratch. This is not supported by this script."
405
+ "You can do it from another script, save it, and load it from here, using --tokenizer_name."
406
+ )
407
+
408
+ # Preprocessing the datasets.
409
+ # First we tokenize all the texts.
410
+ if training_args.do_train:
411
+ column_names = datasets["train"].column_names
412
+ else:
413
+ column_names = datasets["validation"].column_names
414
+ text_column_name = "text" if "text" in column_names else column_names[0]
415
+
416
+ max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)
417
+
418
+ if data_args.line_by_line:
419
+ # When using line_by_line, we just tokenize each nonempty line.
420
+ padding = "max_length" if data_args.pad_to_max_length else False
421
+
422
+ def tokenize_function(examples):
423
+ # Remove empty lines
424
+ examples = [line for line in examples if len(line) > 0 and not line.isspace()]
425
+ return tokenizer(
426
+ examples,
427
+ return_special_tokens_mask=True,
428
+ padding=padding,
429
+ truncation=True,
430
+ max_length=max_seq_length,
431
+ )
432
+
433
+ if data_args.tokenized_dataset_filepath is not None:
434
+ # Loading a tokenized dataset from local file.
435
+ tokenized_datasets = load_from_disk(data_args.tokenized_dataset_filepath)
436
+ else:
437
+ tokenized_datasets = datasets.map(
438
+ tokenize_function,
439
+ input_columns=[text_column_name],
440
+ batched=True,
441
+ num_proc=data_args.preprocessing_num_workers,
442
+ remove_columns=column_names,
443
+ load_from_cache_file=not data_args.overwrite_cache,
444
+ )
445
+
446
+ else:
447
+ # Otherwise, we tokenize every text, then concatenate them together before splitting them in smaller parts.
448
+ # We use `return_special_tokens_mask=True` because DataCollatorForLanguageModeling (see below) is more
449
+ # efficient when it receives the `special_tokens_mask`.
450
+ def tokenize_function(examples):
451
+ return tokenizer(examples[text_column_name], return_special_tokens_mask=True)
452
+
453
+ if data_args.tokenized_dataset_filepath is not None:
454
+ # Loading a tokenized dataset from local file.
455
+ tokenized_datasets = load_from_disk(data_args.tokenized_dataset_filepath)
456
+ else:
457
+ tokenized_datasets = datasets.map(
458
+ tokenize_function,
459
+ batched=True,
460
+ num_proc=data_args.preprocessing_num_workers,
461
+ remove_columns=column_names,
462
+ load_from_cache_file=not data_args.overwrite_cache,
463
+ )
464
+
465
+ # Main data processing function that will concatenate all texts from our dataset and generate chunks of
466
+ # max_seq_length.
467
+ def group_texts(examples):
468
+ # Concatenate all texts.
469
+ concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
470
+ total_length = len(concatenated_examples[list(examples.keys())[0]])
471
+ # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
472
+ # customize this part to your needs.
473
+ if total_length >= max_seq_length:
474
+ total_length = (total_length // max_seq_length) * max_seq_length
475
+ # Split by chunks of max_len.
476
+ result = {
477
+ k: [t[i : i + max_seq_length] for i in range(0, total_length, max_seq_length)]
478
+ for k, t in concatenated_examples.items()
479
+ }
480
+ return result
481
+
482
+ # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a
483
+ # remainder for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value
484
+ # might be slower to preprocess.
485
+ #
486
+ # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
487
+ # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map
488
+ tokenized_datasets = tokenized_datasets.map(
489
+ group_texts,
490
+ batched=True,
491
+ num_proc=data_args.preprocessing_num_workers,
492
+ load_from_cache_file=not data_args.overwrite_cache,
493
+ )
494
+
495
+ # save the tokenized dataset for future runs
496
+ if data_args.save_tokenized_dataset_filepath is not None:
497
+ tokenized_datasets.save_to_disk(data_args.save_tokenized_dataset_filepath)
498
+
499
+ # Enable tensorboard only on the master node
500
+ has_tensorboard = is_tensorboard_available()
501
+ if has_tensorboard and jax.process_index() == 0:
502
+ try:
503
+ from flax.metrics.tensorboard import SummaryWriter
504
+
505
+ summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir))
506
+ except ImportError as ie:
507
+ has_tensorboard = False
508
+ logger.warning(
509
+ f"Unable to display metrics through TensorBoard because some package are not installed: {ie}"
510
+ )
511
+ else:
512
+ logger.warning(
513
+ "Unable to display metrics through TensorBoard because the package is not installed: "
514
+ "Please run pip install tensorboard to enable."
515
+ )
516
+
517
+ # Data collator
518
+ # This one will take care of randomly masking the tokens.
519
+ data_collator = FlaxDataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=data_args.mlm_probability)
520
+
521
+ # Initialize our training
522
+ rng = jax.random.PRNGKey(training_args.seed)
523
+ dropout_rngs = jax.random.split(rng, jax.local_device_count())
524
+
525
+ if model_args.model_name_or_path:
526
+ model = FlaxAutoModelForMaskedLM.from_pretrained(
527
+ model_args.model_name_or_path, config=config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
528
+ )
529
+ else:
530
+ model = FlaxAutoModelForMaskedLM.from_config(
531
+ config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
532
+ )
533
+
534
+ # Store some constant
535
+ num_epochs = int(training_args.num_train_epochs)
536
+ train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
537
+ eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
538
+
539
+ num_train_steps = len(tokenized_datasets["train"]) // train_batch_size * num_epochs
540
+
541
+ # Create learning rate schedule
542
+ warmup_fn = optax.linear_schedule(
543
+ init_value=0.0, end_value=training_args.learning_rate, transition_steps=training_args.warmup_steps
544
+ )
545
+ decay_fn = optax.linear_schedule(
546
+ init_value=training_args.learning_rate,
547
+ end_value=0,
548
+ transition_steps=num_train_steps - training_args.warmup_steps,
549
+ )
550
+ linear_decay_lr_schedule_fn = optax.join_schedules(
551
+ schedules=[warmup_fn, decay_fn], boundaries=[training_args.warmup_steps]
552
+ )
553
+
554
+ # We use Optax's "masking" functionality to not apply weight decay
555
+ # to bias and LayerNorm scale parameters. decay_mask_fn returns a
556
+ # mask boolean with the same structure as the parameters.
557
+ # The mask is True for parameters that should be decayed.
558
+ # Note that this mask is specifically adapted for FlaxBERT-like models.
559
+ # For other models, one should correct the layer norm parameter naming
560
+ # accordingly.
561
+ def decay_mask_fn(params):
562
+ flat_params = traverse_util.flatten_dict(params)
563
+ flat_mask = {path: (path[-1] != "bias" and path[-2:] != ("LayerNorm", "scale")) for path in flat_params}
564
+ return traverse_util.unflatten_dict(flat_mask)
565
+
566
+ # create adam optimizer
567
+ if training_args.adafactor:
568
+ # We use the default parameters here to initialize adafactor,
569
+ # For more details about the parameters please check https://github.com/deepmind/optax/blob/ed02befef9bf81cbbf236be3d2b0e032e9ed4a40/optax/_src/alias.py#L74
570
+ optimizer = optax.adafactor(
571
+ learning_rate=linear_decay_lr_schedule_fn,
572
+ )
573
+ else:
574
+ optimizer = optax.adamw(
575
+ learning_rate=linear_decay_lr_schedule_fn,
576
+ b1=training_args.adam_beta1,
577
+ b2=training_args.adam_beta2,
578
+ eps=training_args.adam_epsilon,
579
+ weight_decay=training_args.weight_decay,
580
+ mask=decay_mask_fn,
581
+ )
582
+
583
+ # Setup train state
584
+ state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=optimizer)
585
+
586
+ # Define gradient update step fn
587
+ def train_step(state, batch, dropout_rng):
588
+ dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)
589
+
590
+ def loss_fn(params):
591
+ labels = batch.pop("labels")
592
+
593
+ logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
594
+
595
+ # compute loss, ignore padded input tokens
596
+ label_mask = jnp.where(labels > 0, 1.0, 0.0)
597
+ loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])) * label_mask
598
+
599
+ # take average
600
+ loss = loss.sum() / label_mask.sum()
601
+
602
+ return loss
603
+
604
+ grad_fn = jax.value_and_grad(loss_fn)
605
+ loss, grad = grad_fn(state.params)
606
+ grad = jax.lax.pmean(grad, "batch")
607
+ new_state = state.apply_gradients(grads=grad)
608
+
609
+ metrics = jax.lax.pmean(
610
+ {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}, axis_name="batch"
611
+ )
612
+
613
+ return new_state, metrics, new_dropout_rng
614
+
615
+ # Create parallel version of the train step
616
+ p_train_step = jax.pmap(train_step, "batch", donate_argnums=(0,))
617
+
618
+ # Define eval fn
619
+ def eval_step(params, batch):
620
+ labels = batch.pop("labels")
621
+
622
+ logits = model(**batch, params=params, train=False)[0]
623
+
624
+ # compute loss, ignore padded input tokens
625
+ label_mask = jnp.where(labels > 0, 1.0, 0.0)
626
+ loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])) * label_mask
627
+
628
+ # compute accuracy
629
+ accuracy = jnp.equal(jnp.argmax(logits, axis=-1), labels) * label_mask
630
+
631
+ # summarize metrics
632
+ metrics = {"loss": loss.sum(), "accuracy": accuracy.sum(), "normalizer": label_mask.sum()}
633
+ metrics = jax.lax.psum(metrics, axis_name="batch")
634
+
635
+ return metrics
636
+
637
+ p_eval_step = jax.pmap(eval_step, "batch", donate_argnums=(0,))
638
+
639
+ # Replicate the train state on each device
640
+ state = jax_utils.replicate(state)
641
+
642
+ train_time = 0
643
+ epochs = tqdm(range(num_epochs), desc=f"Epoch ... (1/{num_epochs})", position=0)
644
+ for epoch in epochs:
645
+ # ======================== Training ================================
646
+ train_start = time.time()
647
+ train_metrics = []
648
+
649
+ # Create sampling rng
650
+ rng, input_rng = jax.random.split(rng)
651
+
652
+ # Generate an epoch by shuffling sampling indices from the train dataset
653
+ num_train_samples = len(tokenized_datasets["train"])
654
+ train_samples_idx = jax.random.permutation(input_rng, jnp.arange(num_train_samples))
655
+ train_batch_idx = generate_batch_splits(train_samples_idx, train_batch_size)
656
+
657
+ # Gather the indexes for creating the batch and do a training step
658
+ for step, batch_idx in enumerate(tqdm(train_batch_idx, desc="Training...", position=1)):
659
+ samples = [tokenized_datasets["train"][int(idx)] for idx in batch_idx]
660
+ model_inputs = data_collator(samples, pad_to_multiple_of=16)
661
+
662
+ # Model forward
663
+ model_inputs = shard(model_inputs.data)
664
+ state, train_metric, dropout_rngs = p_train_step(state, model_inputs, dropout_rngs)
665
+ train_metrics.append(train_metric)
666
+
667
+ cur_step = epoch * (num_train_samples // train_batch_size) + step
668
+
669
+ if cur_step % training_args.logging_steps == 0 and cur_step > 0:
670
+ # Save metrics
671
+ train_metric = jax_utils.unreplicate(train_metric)
672
+ train_time += time.time() - train_start
673
+ if has_tensorboard and jax.process_index() == 0:
674
+ write_train_metric(summary_writer, train_metrics, train_time, cur_step)
675
+
676
+ epochs.write(
677
+ f"Step... ({cur_step} | Loss: {train_metric['loss']}, Learning Rate: {train_metric['learning_rate']})"
678
+ )
679
+
680
+ train_metrics = []
681
+
682
+ if cur_step % training_args.eval_steps == 0 and cur_step > 0:
683
+ # ======================== Evaluating ==============================
684
+ num_eval_samples = len(tokenized_datasets["validation"])
685
+ eval_samples_idx = jnp.arange(num_eval_samples)
686
+ eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size)
687
+
688
+ eval_metrics = []
689
+ for i, batch_idx in enumerate(tqdm(eval_batch_idx, desc="Evaluating ...", position=2)):
690
+ samples = [tokenized_datasets["validation"][int(idx)] for idx in batch_idx]
691
+ model_inputs = data_collator(samples, pad_to_multiple_of=16)
692
+
693
+ # Model forward
694
+ model_inputs = shard(model_inputs.data)
695
+ metrics = p_eval_step(state.params, model_inputs)
696
+ eval_metrics.append(metrics)
697
+
698
+ # normalize eval metrics
699
+ eval_metrics = get_metrics(eval_metrics)
700
+ eval_metrics = jax.tree_map(jnp.sum, eval_metrics)
701
+ eval_normalizer = eval_metrics.pop("normalizer")
702
+ eval_metrics = jax.tree_map(lambda x: x / eval_normalizer, eval_metrics)
703
+
704
+ # Update progress bar
705
+ epochs.desc = f"Step... ({cur_step} | Loss: {eval_metrics['loss']}, Acc: {eval_metrics['accuracy']})"
706
+
707
+ # Save metrics
708
+ if has_tensorboard and jax.process_index() == 0:
709
+ write_eval_metric(summary_writer, eval_metrics, cur_step)
710
+
711
+ if cur_step % training_args.save_steps == 0 and cur_step > 0:
712
+ # save checkpoint after each epoch and push checkpoint to the hub
713
+ if jax.process_index() == 0:
714
+ params = jax.device_get(jax.tree_map(lambda x: x[0], state.params))
715
+ model.save_pretrained(
716
+ training_args.output_dir,
717
+ params=params,
718
+ push_to_hub=training_args.push_to_hub,
719
+ commit_message=f"Saving weights and logs of step {cur_step}",
720
+ )
721
+
722
+ # save also at the end of epoch
723
+ try:
724
+ if jax.process_index() == 0:
725
+ params = jax.device_get(jax.tree_map(lambda x: x[0], state.params))
726
+ model.save_pretrained(
727
+ training_args.output_dir,
728
+ params=params,
729
+ push_to_hub=training_args.push_to_hub,
730
+ commit_message=f"Saving weights and logs of epoch {epoch}",
731
+ )
732
+ except:
733
+ # push to hub fails the whole script if nothing new to commit
734
+ pass
special_tokens_map.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"bos_token": "<s>", "eos_token": "</s>", "unk_token": "<unk>", "sep_token": "</s>", "pad_token": "<pad>", "cls_token": "<s>", "mask_token": {"content": "<mask>", "single_word": false, "lstrip": true, "rstrip": false, "normalized": false}}
start_train.sh ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ unset LD_PRELOAD
2
+ python3 run_mlm_flax.py \
3
+ --output_dir="./" \
4
+ --model_type="roberta" \
5
+ --config_name="./" \
6
+ --tokenizer_name="./" \
7
+ --dataset_filepath="/researchdisk1/data/training_data_full" \
8
+ --save_tokenized_dataset_filepath="/researchdisk1/data/training_data_full_tokenized_128" \
9
+ --max_seq_length="128" \
10
+ --pad_to_max_length \
11
+ --preprocessing_num_workers="96" \
12
+ --per_device_train_batch_size="128" \
13
+ --per_device_eval_batch_size="128" \
14
+ --learning_rate="3e-4" \
15
+ --warmup_steps="1000" \
16
+ --overwrite_output_dir \
17
+ --num_train_epochs="3" \
18
+ --save_strategy="steps" \
19
+ --save_steps="10000" \
20
+ --save_total_limit="5" \
21
+ --eval_steps="10000" \
22
+ --logging_steps="1000" \
23
+ --dtype="bfloat16" \
24
+ --push_to_hub
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"unk_token": "<unk>", "bos_token": "<s>", "eos_token": "</s>", "add_prefix_space": false, "errors": "replace", "sep_token": "</s>", "cls_token": "<s>", "pad_token": "<pad>", "mask_token": "<mask>", "special_tokens_map_file": null, "name_or_path": "./", "tokenizer_class": "RobertaTokenizer"}
train_tokenizer.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from datasets import load_dataset, load_from_disk
2
+ from tokenizers import trainers, Tokenizer, normalizers, ByteLevelBPETokenizer
3
+ from transformers import AutoConfig, AutoTokenizer
4
+
5
+
6
+ model_dir = "./" # ${MODEL_DIR}
7
+
8
+ # load roberta-large config
9
+ config = AutoConfig.from_pretrained("roberta-large")
10
+ config.save_pretrained(model_dir)
11
+
12
+ # load dataset
13
+ dataset = load_from_disk("/researchdisk1/data/training_data_full")
14
+ dataset = dataset["train"]
15
+
16
+ # Instantiate tokenizer
17
+ tokenizer = ByteLevelBPETokenizer()
18
+ def batch_iterator(batch_size=1000):
19
+ for i in range(0, len(dataset), batch_size):
20
+ yield dataset[i: i + batch_size]["text"]
21
+
22
+ # Customized training
23
+ tokenizer.train_from_iterator(batch_iterator(), vocab_size=config.vocab_size, min_frequency=2, special_tokens=[
24
+ "<s>",
25
+ "<pad>",
26
+ "</s>",
27
+ "<unk>",
28
+ "<mask>",
29
+ ])
30
+ # Save files to disk
31
+ tokenizer.save(f"{model_dir}/tokenizer.json")
32
+
33
+ tokenizer = AutoTokenizer.from_pretrained(model_dir)
34
+ tokenizer.save_pretrained(model_dir)
vocab.json ADDED
The diff for this file is too large to render. See raw diff