pere commited on
Commit
d761f2e
1 Parent(s): be52aca

first push

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