run_mlm.py 29.2 KB
Newer Older
1
#!/usr/bin/env python
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
# coding=utf-8
# Copyright 2020 The HuggingFace Team All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Fine-tuning the library models for masked language modeling (BERT, ALBERT, RoBERTa...) on a text file or a dataset.

Here is the full list of checkpoints on the hub that can be fine-tuned by this script:
20
https://huggingface.co/models?filter=fill-mask
21
22
23
24
25
26
27
28
"""
# You can also adapt this script on your own masked language modeling task. Pointers for this are left as comments.

import logging
import math
import os
import sys
from dataclasses import dataclass, field
29
from itertools import chain
30
31
from typing import Optional

32
import datasets
33
import evaluate
34
import torch
35
from datasets import load_dataset
36
37
38
39
40
41
42
43
44
45
46
47

import transformers
from transformers import (
    CONFIG_MAPPING,
    MODEL_FOR_MASKED_LM_MAPPING,
    AutoConfig,
    AutoModelForMaskedLM,
    AutoTokenizer,
    DataCollatorForLanguageModeling,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
48
    is_torch_xla_available,
49
50
    set_seed,
)
51
from transformers.trainer_utils import get_last_checkpoint
52
from transformers.utils import check_min_version, send_example_telemetry
53
from transformers.utils.versions import require_version
54
55


56
# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
Arthur Zucker's avatar
Arthur Zucker committed
57
check_min_version("4.42.0.dev0")
Sylvain Gugger's avatar
Sylvain Gugger committed
58

59
require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt")
60

61
62
63
64
65
66
67
68
69
70
71
72
73
74
logger = logging.getLogger(__name__)
MODEL_CONFIG_CLASSES = list(MODEL_FOR_MASKED_LM_MAPPING.keys())
MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)


@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
    """

    model_name_or_path: Optional[str] = field(
        default=None,
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
75
            "help": (
76
                "The model checkpoint for weights initialization. Don't set if you want to train a model from scratch."
Sylvain Gugger's avatar
Sylvain Gugger committed
77
            )
78
79
80
81
82
83
        },
    )
    model_type: Optional[str] = field(
        default=None,
        metadata={"help": "If training from scratch, pass a model type from the list: " + ", ".join(MODEL_TYPES)},
    )
84
85
86
    config_overrides: Optional[str] = field(
        default=None,
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
87
88
89
90
            "help": (
                "Override some existing default config settings when a model is trained from scratch. Example: "
                "n_embd=10,resid_pdrop=0.2,scale_attn_weights=false,summary_type=cls_index"
            )
91
92
        },
    )
93
94
95
96
97
98
99
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
    )
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
    )
    cache_dir: Optional[str] = field(
100
101
        default=None,
        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
102
103
104
105
106
    )
    use_fast_tokenizer: bool = field(
        default=True,
        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
    )
107
108
109
110
    model_revision: str = field(
        default="main",
        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
    )
111
112
    token: str = field(
        default=None,
113
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
114
            "help": (
115
116
                "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token "
                "generated when running `huggingface-cli login` (stored in `~/.huggingface`)."
Sylvain Gugger's avatar
Sylvain Gugger committed
117
            )
118
119
        },
    )
120
121
122
123
    trust_remote_code: bool = field(
        default=False,
        metadata={
            "help": (
124
                "Whether or not to allow for custom models defined on the Hub in their own modeling files. This option "
125
                "should only be set to `True` for repositories you trust and in which you have read the code, as it will "
126
127
128
129
                "execute code present on the Hub on your local machine."
            )
        },
    )
130
131
132
133
134
135
136
137
138
139
    torch_dtype: Optional[str] = field(
        default=None,
        metadata={
            "help": (
                "Override the default `torch.dtype` and load the model under this dtype. If `auto` is passed, the "
                "dtype will be automatically derived from the model's weights."
            ),
            "choices": ["auto", "bfloat16", "float16", "float32"],
        },
    )
140
141
142
143
    low_cpu_mem_usage: bool = field(
        default=False,
        metadata={
            "help": (
144
                "It is an option to create the model as an empty shell, then only materialize its parameters when the pretrained weights are loaded. "
145
146
147
148
                "set True will benefit LLM loading time and RAM consumption."
            )
        },
    )
149

150
151
152
153
154
155
    def __post_init__(self):
        if self.config_overrides is not None and (self.config_name is not None or self.model_name_or_path is not None):
            raise ValueError(
                "--config_overrides can't be used in combination with --config_name or --model_name_or_path"
            )

156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176

@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    """

    dataset_name: Optional[str] = field(
        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
    )
    dataset_config_name: Optional[str] = field(
        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
    )
    train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
    validation_file: Optional[str] = field(
        default=None,
        metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},
    )
    overwrite_cache: bool = field(
        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
    )
177
178
179
180
181
182
    validation_split_percentage: Optional[int] = field(
        default=5,
        metadata={
            "help": "The percentage of the train set used as validation set in case there's no validation split"
        },
    )
183
184
185
    max_seq_length: Optional[int] = field(
        default=None,
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
186
187
188
189
            "help": (
                "The maximum total input sequence length after tokenization. Sequences longer "
                "than this will be truncated."
            )
190
191
192
193
194
195
196
197
198
        },
    )
    preprocessing_num_workers: Optional[int] = field(
        default=None,
        metadata={"help": "The number of processes to use for the preprocessing."},
    )
    mlm_probability: float = field(
        default=0.15, metadata={"help": "Ratio of tokens to mask for masked language modeling loss"}
    )
199
200
201
202
203
204
205
    line_by_line: bool = field(
        default=False,
        metadata={"help": "Whether distinct lines of text in the dataset are to be handled as distinct sequences."},
    )
    pad_to_max_length: bool = field(
        default=False,
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
206
207
208
209
            "help": (
                "Whether to pad all samples to `max_seq_length`. "
                "If False, will pad the samples dynamically when batching to the maximum length in the batch."
            )
210
211
        },
    )
212
213
214
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
215
216
217
218
            "help": (
                "For debugging purposes or quicker training, truncate the number of training examples to this "
                "value if set."
            )
219
220
        },
    )
221
    max_eval_samples: Optional[int] = field(
222
223
        default=None,
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
224
225
226
227
            "help": (
                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
                "value if set."
            )
228
229
        },
    )
230
    streaming: bool = field(default=False, metadata={"help": "Enable streaming mode"})
231
232

    def __post_init__(self):
233
234
235
        if self.streaming:
            require_version("datasets>=2.0.0", "The streaming feature requires `datasets>=2.0.0`")

236
237
238
239
240
        if self.dataset_name is None and self.train_file is None and self.validation_file is None:
            raise ValueError("Need either a dataset name or a training/validation file.")
        else:
            if self.train_file is not None:
                extension = self.train_file.split(".")[-1]
241
242
                if extension not in ["csv", "json", "txt"]:
                    raise ValueError("`train_file` should be a csv, a json or a txt file.")
243
244
            if self.validation_file is not None:
                extension = self.validation_file.split(".")[-1]
245
246
                if extension not in ["csv", "json", "txt"]:
                    raise ValueError("`validation_file` should be a csv, a json or a txt file.")
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261


def main():
    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

262
263
264
265
    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
    # information sent is the one passed as arguments along with your Python/PyTorch versions.
    send_example_telemetry("run_mlm", model_args, data_args)

266
267
    # Setup logging
    logging.basicConfig(
268
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
269
        datefmt="%m/%d/%Y %H:%M:%S",
270
        handlers=[logging.StreamHandler(sys.stdout)],
271
    )
272

273
274
275
276
    if training_args.should_log:
        # The default of training_args.log_level is passive, so we set log level at info here to have that default.
        transformers.utils.logging.set_verbosity_info()

277
278
279
280
281
282
    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()
283
284
285

    # Log on each process the small summary:
    logger.warning(
286
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
287
        + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}"
288
289
    )
    # Set the verbosity to info of the Transformers logger (on main process only):
290
    logger.info(f"Training/evaluation parameters {training_args}")
291

292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

307
308
309
310
311
312
313
314
315
316
317
318
319
320
    # Set seed before initializing model.
    set_seed(training_args.seed)

    # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
    # (the dataset will be downloaded automatically from the datasets Hub
    #
    # For CSV/JSON files, this script will use the column called 'text' or the first column. You can easily tweak this
    # behavior (see below)
    #
    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
    # download the dataset.
    if data_args.dataset_name is not None:
        # Downloading and loading a dataset from the hub.
321
        raw_datasets = load_dataset(
322
323
324
            data_args.dataset_name,
            data_args.dataset_config_name,
            cache_dir=model_args.cache_dir,
325
            token=model_args.token,
326
            streaming=data_args.streaming,
327
328
329
        )
        if "validation" not in raw_datasets.keys():
            raw_datasets["validation"] = load_dataset(
330
331
332
                data_args.dataset_name,
                data_args.dataset_config_name,
                split=f"train[:{data_args.validation_split_percentage}%]",
333
                cache_dir=model_args.cache_dir,
334
                token=model_args.token,
335
                streaming=data_args.streaming,
336
            )
337
            raw_datasets["train"] = load_dataset(
338
339
340
                data_args.dataset_name,
                data_args.dataset_config_name,
                split=f"train[{data_args.validation_split_percentage}%:]",
341
                cache_dir=model_args.cache_dir,
342
                token=model_args.token,
343
                streaming=data_args.streaming,
344
            )
345
346
347
348
    else:
        data_files = {}
        if data_args.train_file is not None:
            data_files["train"] = data_args.train_file
349
            extension = data_args.train_file.split(".")[-1]
350
        if data_args.validation_file is not None:
351
            data_files["validation"] = data_args.validation_file
352
            extension = data_args.validation_file.split(".")[-1]
353
354
        if extension == "txt":
            extension = "text"
355
356
357
358
        raw_datasets = load_dataset(
            extension,
            data_files=data_files,
            cache_dir=model_args.cache_dir,
359
            token=model_args.token,
360
        )
361
362
363
364
365
366
367
368

        # If no validation data is there, validation_split_percentage will be used to divide the dataset.
        if "validation" not in raw_datasets.keys():
            raw_datasets["validation"] = load_dataset(
                extension,
                data_files=data_files,
                split=f"train[:{data_args.validation_split_percentage}%]",
                cache_dir=model_args.cache_dir,
369
                token=model_args.token,
370
371
372
373
374
375
            )
            raw_datasets["train"] = load_dataset(
                extension,
                data_files=data_files,
                split=f"train[{data_args.validation_split_percentage}%:]",
                cache_dir=model_args.cache_dir,
376
                token=model_args.token,
377
378
            )

379
    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
380
    # https://huggingface.co/docs/datasets/loading_datasets.
381
382
383
384
385
386

    # Load pretrained model and tokenizer
    #
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.
387
388
389
    config_kwargs = {
        "cache_dir": model_args.cache_dir,
        "revision": model_args.model_revision,
390
        "token": model_args.token,
391
        "trust_remote_code": model_args.trust_remote_code,
392
    }
393
    if model_args.config_name:
394
        config = AutoConfig.from_pretrained(model_args.config_name, **config_kwargs)
395
    elif model_args.model_name_or_path:
396
        config = AutoConfig.from_pretrained(model_args.model_name_or_path, **config_kwargs)
397
398
399
    else:
        config = CONFIG_MAPPING[model_args.model_type]()
        logger.warning("You are instantiating a new config instance from scratch.")
400
401
402
        if model_args.config_overrides is not None:
            logger.info(f"Overriding config: {model_args.config_overrides}")
            config.update_from_string(model_args.config_overrides)
403
            logger.info(f"New config: {config}")
404

405
406
407
408
    tokenizer_kwargs = {
        "cache_dir": model_args.cache_dir,
        "use_fast": model_args.use_fast_tokenizer,
        "revision": model_args.model_revision,
409
        "token": model_args.token,
410
        "trust_remote_code": model_args.trust_remote_code,
411
    }
412
    if model_args.tokenizer_name:
413
        tokenizer = AutoTokenizer.from_pretrained(model_args.tokenizer_name, **tokenizer_kwargs)
414
    elif model_args.model_name_or_path:
415
        tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, **tokenizer_kwargs)
416
417
    else:
        raise ValueError(
418
            "You are instantiating a new tokenizer from scratch. This is not supported by this script. "
419
420
421
422
            "You can do it from another script, save it, and load it from here, using --tokenizer_name."
        )

    if model_args.model_name_or_path:
423
424
425
426
427
        torch_dtype = (
            model_args.torch_dtype
            if model_args.torch_dtype in ["auto", None]
            else getattr(torch, model_args.torch_dtype)
        )
428
429
430
431
432
        model = AutoModelForMaskedLM.from_pretrained(
            model_args.model_name_or_path,
            from_tf=bool(".ckpt" in model_args.model_name_or_path),
            config=config,
            cache_dir=model_args.cache_dir,
433
            revision=model_args.model_revision,
434
            token=model_args.token,
435
            trust_remote_code=model_args.trust_remote_code,
436
            torch_dtype=torch_dtype,
437
            low_cpu_mem_usage=model_args.low_cpu_mem_usage,
438
439
440
        )
    else:
        logger.info("Training new model from scratch")
441
        model = AutoModelForMaskedLM.from_config(config, trust_remote_code=model_args.trust_remote_code)
442

443
444
445
446
447
    # We resize the embeddings only when necessary to avoid index errors. If you are creating a model from scratch
    # on a small vocab and want a smaller embedding size, remove this test.
    embedding_size = model.get_input_embeddings().weight.shape[0]
    if len(tokenizer) > embedding_size:
        model.resize_token_embeddings(len(tokenizer))
448
449
450
451

    # Preprocessing the datasets.
    # First we tokenize all the texts.
    if training_args.do_train:
452
        column_names = list(raw_datasets["train"].features)
453
    else:
454
        column_names = list(raw_datasets["validation"].features)
455
456
    text_column_name = "text" if "text" in column_names else column_names[0]

457
458
459
    if data_args.max_seq_length is None:
        max_seq_length = tokenizer.model_max_length
        if max_seq_length > 1024:
460
            logger.warning(
461
462
463
                "The chosen tokenizer supports a `model_max_length` that is longer than the default `block_size` value"
                " of 1024. If you would like to use a longer `block_size` up to `tokenizer.model_max_length` you can"
                " override this default with `--block_size xxx`."
464
465
466
467
            )
            max_seq_length = 1024
    else:
        if data_args.max_seq_length > tokenizer.model_max_length:
468
            logger.warning(
469
                f"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the "
470
471
472
473
                f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
            )
        max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)

474
475
476
477
478
479
    if data_args.line_by_line:
        # When using line_by_line, we just tokenize each nonempty line.
        padding = "max_length" if data_args.pad_to_max_length else False

        def tokenize_function(examples):
            # Remove empty lines
480
481
482
            examples[text_column_name] = [
                line for line in examples[text_column_name] if len(line) > 0 and not line.isspace()
            ]
483
            return tokenizer(
484
                examples[text_column_name],
485
486
                padding=padding,
                truncation=True,
487
                max_length=max_seq_length,
488
489
490
491
                # We use this option because DataCollatorForLanguageModeling (see below) is more efficient when it
                # receives the `special_tokens_mask`.
                return_special_tokens_mask=True,
            )
492

493
        with training_args.main_process_first(desc="dataset map tokenization"):
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
            if not data_args.streaming:
                tokenized_datasets = raw_datasets.map(
                    tokenize_function,
                    batched=True,
                    num_proc=data_args.preprocessing_num_workers,
                    remove_columns=[text_column_name],
                    load_from_cache_file=not data_args.overwrite_cache,
                    desc="Running tokenizer on dataset line_by_line",
                )
            else:
                tokenized_datasets = raw_datasets.map(
                    tokenize_function,
                    batched=True,
                    remove_columns=[text_column_name],
                )
509
510
    else:
        # Otherwise, we tokenize every text, then concatenate them together before splitting them in smaller parts.
511
512
        # We use `return_special_tokens_mask=True` because DataCollatorForLanguageModeling (see below) is more
        # efficient when it receives the `special_tokens_mask`.
513
        def tokenize_function(examples):
514
            return tokenizer(examples[text_column_name], return_special_tokens_mask=True)
515

516
        with training_args.main_process_first(desc="dataset map tokenization"):
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
            if not data_args.streaming:
                tokenized_datasets = raw_datasets.map(
                    tokenize_function,
                    batched=True,
                    num_proc=data_args.preprocessing_num_workers,
                    remove_columns=column_names,
                    load_from_cache_file=not data_args.overwrite_cache,
                    desc="Running tokenizer on every text in dataset",
                )
            else:
                tokenized_datasets = raw_datasets.map(
                    tokenize_function,
                    batched=True,
                    remove_columns=column_names,
                )
532
533
534
535
536

        # Main data processing function that will concatenate all texts from our dataset and generate chunks of
        # max_seq_length.
        def group_texts(examples):
            # Concatenate all texts.
537
            concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
538
            total_length = len(concatenated_examples[list(examples.keys())[0]])
539
540
541
            # We drop the small remainder, and if the total_length < max_seq_length  we exclude this batch and return an empty dict.
            # We could add padding if the model supported it instead of this drop, you can customize this part to your needs.
            total_length = (total_length // max_seq_length) * max_seq_length
542
543
544
545
546
547
548
549
550
551
552
553
            # Split by chunks of max_len.
            result = {
                k: [t[i : i + max_seq_length] for i in range(0, total_length, max_seq_length)]
                for k, t in concatenated_examples.items()
            }
            return result

        # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a
        # remainder for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value
        # might be slower to preprocess.
        #
        # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
554
        # https://huggingface.co/docs/datasets/process#map
555

556
        with training_args.main_process_first(desc="grouping texts together"):
557
558
559
560
561
562
563
564
565
566
567
568
569
            if not data_args.streaming:
                tokenized_datasets = tokenized_datasets.map(
                    group_texts,
                    batched=True,
                    num_proc=data_args.preprocessing_num_workers,
                    load_from_cache_file=not data_args.overwrite_cache,
                    desc=f"Grouping texts in chunks of {max_seq_length}",
                )
            else:
                tokenized_datasets = tokenized_datasets.map(
                    group_texts,
                    batched=True,
                )
570

571
572
573
574
575
    if training_args.do_train:
        if "train" not in tokenized_datasets:
            raise ValueError("--do_train requires a train dataset")
        train_dataset = tokenized_datasets["train"]
        if data_args.max_train_samples is not None:
576
577
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
            train_dataset = train_dataset.select(range(max_train_samples))
578
579
580
581
582

    if training_args.do_eval:
        if "validation" not in tokenized_datasets:
            raise ValueError("--do_eval requires a validation dataset")
        eval_dataset = tokenized_datasets["validation"]
583
        if data_args.max_eval_samples is not None:
584
585
            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
            eval_dataset = eval_dataset.select(range(max_eval_samples))
586

587
        def preprocess_logits_for_metrics(logits, labels):
davidleonfdez's avatar
davidleonfdez committed
588
589
590
591
            if isinstance(logits, tuple):
                # Depending on the model and config, logits may contain extra tensors,
                # like past_key_values, but logits always come first
                logits = logits[0]
592
593
            return logits.argmax(dim=-1)

594
        metric = evaluate.load("accuracy", cache_dir=model_args.cache_dir)
595
596
597
598
599
600
601
602
603
604
605
606

        def compute_metrics(eval_preds):
            preds, labels = eval_preds
            # preds have the same shape as the labels, after the argmax(-1) has been calculated
            # by preprocess_logits_for_metrics
            labels = labels.reshape(-1)
            preds = preds.reshape(-1)
            mask = labels != -100
            labels = labels[mask]
            preds = preds[mask]
            return metric.compute(predictions=preds, references=labels)

607
608
    # Data collator
    # This one will take care of randomly masking the tokens.
609
610
611
612
613
614
    pad_to_multiple_of_8 = data_args.line_by_line and training_args.fp16 and not data_args.pad_to_max_length
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm_probability=data_args.mlm_probability,
        pad_to_multiple_of=8 if pad_to_multiple_of_8 else None,
    )
615
616
617
618
619

    # Initialize our Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
620
621
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
622
623
        tokenizer=tokenizer,
        data_collator=data_collator,
624
        compute_metrics=compute_metrics if training_args.do_eval and not is_torch_xla_available() else None,
625
        preprocess_logits_for_metrics=preprocess_logits_for_metrics
626
        if training_args.do_eval and not is_torch_xla_available()
627
        else None,
628
629
630
631
    )

    # Training
    if training_args.do_train:
632
633
634
635
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
636
637
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
638
        trainer.save_model()  # Saves the tokenizer too for easy upload
639
        metrics = train_result.metrics
640

641
642
643
644
645
        max_train_samples = (
            data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
        )
        metrics["train_samples"] = min(max_train_samples, len(train_dataset))

646
647
648
        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()
649

650
651
652
653
    # Evaluation
    if training_args.do_eval:
        logger.info("*** Evaluate ***")

654
        metrics = trainer.evaluate()
655

656
657
        max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)
        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))
658
659
660
661
        try:
            perplexity = math.exp(metrics["eval_loss"])
        except OverflowError:
            perplexity = float("inf")
662
        metrics["perplexity"] = perplexity
663

664
665
        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)
666

667
668
669
670
671
672
673
674
    kwargs = {"finetuned_from": model_args.model_name_or_path, "tasks": "fill-mask"}
    if data_args.dataset_name is not None:
        kwargs["dataset_tags"] = data_args.dataset_name
        if data_args.dataset_config_name is not None:
            kwargs["dataset_args"] = data_args.dataset_config_name
            kwargs["dataset"] = f"{data_args.dataset_name} {data_args.dataset_config_name}"
        else:
            kwargs["dataset"] = data_args.dataset_name
Sylvain Gugger's avatar
Sylvain Gugger committed
675

676
    if training_args.push_to_hub:
Sylvain Gugger's avatar
Sylvain Gugger committed
677
        trainer.push_to_hub(**kwargs)
678
679
    else:
        trainer.create_model_card(**kwargs)
Sylvain Gugger's avatar
Sylvain Gugger committed
680

681
682
683
684
685
686
687
688

def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()


if __name__ == "__main__":
    main()