run_qa.py 32.2 KB
Newer Older
1
#!/usr/bin/env python
Sylvain Gugger's avatar
Sylvain Gugger committed
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# 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.
"""
17
Fine-tuning the library models for question answering using a slightly adapted version of the 馃 Trainer.
Sylvain Gugger's avatar
Sylvain Gugger committed
18
19
20
21
22
23
"""
# You can also adapt this script on your own question answering task. Pointers for this are left as comments.

import logging
import os
import sys
24
import warnings
Sylvain Gugger's avatar
Sylvain Gugger committed
25
26
27
from dataclasses import dataclass, field
from typing import Optional

28
import datasets
29
import evaluate
30
from datasets import load_dataset
31
32
from trainer_qa import QuestionAnsweringTrainer
from utils_qa import postprocess_qa_predictions
Sylvain Gugger's avatar
Sylvain Gugger committed
33
34
35
36
37
38
39
40
41
42
43
44
45
46

import transformers
from transformers import (
    AutoConfig,
    AutoModelForQuestionAnswering,
    AutoTokenizer,
    DataCollatorWithPadding,
    EvalPrediction,
    HfArgumentParser,
    PreTrainedTokenizerFast,
    TrainingArguments,
    default_data_collator,
    set_seed,
)
47
from transformers.trainer_utils import get_last_checkpoint
48
from transformers.utils import check_min_version, send_example_telemetry
49
from transformers.utils.versions import require_version
Sylvain Gugger's avatar
Sylvain Gugger committed
50
51


52
# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
Arthur Zucker's avatar
Arthur Zucker committed
53
check_min_version("4.45.0.dev0")
Sylvain Gugger's avatar
Sylvain Gugger committed
54

55
require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/question-answering/requirements.txt")
56

Sylvain Gugger's avatar
Sylvain Gugger committed
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
logger = logging.getLogger(__name__)


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

    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
    )
    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(
        default=None,
        metadata={"help": "Path to directory to store the pretrained models downloaded from huggingface.co"},
    )
79
80
81
82
    model_revision: str = field(
        default="main",
        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
    )
83
84
    token: str = field(
        default=None,
85
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
86
            "help": (
87
88
                "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
89
            )
90
91
        },
    )
92
93
94
95
    trust_remote_code: bool = field(
        default=False,
        metadata={
            "help": (
96
97
98
                "Whether to trust the execution of code from datasets/models defined on the Hub."
                " This option should only be set to `True` for repositories you trust and in which you have read the"
                " code, as it will execute code present on the Hub on your local machine."
99
100
101
            )
        },
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120


@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)."},
    )
121
122
123
124
    test_file: Optional[str] = field(
        default=None,
        metadata={"help": "An optional input test data file to evaluate the perplexity on (a text file)."},
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
125
126
127
128
129
130
131
132
133
134
    overwrite_cache: bool = field(
        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
    )
    preprocessing_num_workers: Optional[int] = field(
        default=None,
        metadata={"help": "The number of processes to use for the preprocessing."},
    )
    max_seq_length: int = field(
        default=384,
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
135
136
137
138
            "help": (
                "The maximum total input sequence length after tokenization. Sequences longer "
                "than this will be truncated, sequences shorter will be padded."
            )
Sylvain Gugger's avatar
Sylvain Gugger committed
139
140
141
142
143
        },
    )
    pad_to_max_length: bool = field(
        default=True,
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
144
145
146
147
            "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 (which can be faster on GPU but will be slower on TPU)."
            )
Sylvain Gugger's avatar
Sylvain Gugger committed
148
149
        },
    )
150
151
152
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
153
154
155
156
            "help": (
                "For debugging purposes or quicker training, truncate the number of training examples to this "
                "value if set."
            )
157
158
        },
    )
159
    max_eval_samples: Optional[int] = field(
160
161
        default=None,
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
162
163
164
165
            "help": (
                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
                "value if set."
            )
166
167
        },
    )
168
    max_predict_samples: Optional[int] = field(
169
170
        default=None,
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
171
172
173
174
            "help": (
                "For debugging purposes or quicker training, truncate the number of prediction examples to this "
                "value if set."
            )
175
176
        },
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
177
178
179
180
181
182
    version_2_with_negative: bool = field(
        default=False, metadata={"help": "If true, some of the examples do not have an answer."}
    )
    null_score_diff_threshold: float = field(
        default=0.0,
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
183
184
185
186
187
            "help": (
                "The threshold used to select the null answer: if the best answer has a score that is less than "
                "the score of the null answer minus this threshold, the null answer is selected for this example. "
                "Only useful when `version_2_with_negative=True`."
            )
Sylvain Gugger's avatar
Sylvain Gugger committed
188
189
190
191
192
193
194
195
196
197
198
199
200
        },
    )
    doc_stride: int = field(
        default=128,
        metadata={"help": "When splitting up a long document into chunks, how much stride to take between chunks."},
    )
    n_best_size: int = field(
        default=20,
        metadata={"help": "The total number of n-best predictions to generate when looking for an answer."},
    )
    max_answer_length: int = field(
        default=30,
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
201
202
203
204
            "help": (
                "The maximum length of an answer that can be generated. This is needed because the start "
                "and end predictions are not conditioned on one another."
            )
Sylvain Gugger's avatar
Sylvain Gugger committed
205
206
207
208
        },
    )

    def __post_init__(self):
209
210
211
212
213
214
215
        if (
            self.dataset_name is None
            and self.train_file is None
            and self.validation_file is None
            and self.test_file is None
        ):
            raise ValueError("Need either a dataset name or a training/validation file/test_file.")
Sylvain Gugger's avatar
Sylvain Gugger committed
216
217
218
219
220
221
222
        else:
            if self.train_file is not None:
                extension = self.train_file.split(".")[-1]
                assert extension in ["csv", "json"], "`train_file` should be a csv or a json file."
            if self.validation_file is not None:
                extension = self.validation_file.split(".")[-1]
                assert extension in ["csv", "json"], "`validation_file` should be a csv or a json file."
223
224
225
            if self.test_file is not None:
                extension = self.test_file.split(".")[-1]
                assert extension in ["csv", "json"], "`test_file` should be a csv or a json file."
Sylvain Gugger's avatar
Sylvain Gugger committed
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240


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()

241
242
243
244
    # 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_qa", model_args, data_args)

Sylvain Gugger's avatar
Sylvain Gugger committed
245
246
    # Setup logging
    logging.basicConfig(
247
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
Sylvain Gugger's avatar
Sylvain Gugger committed
248
        datefmt="%m/%d/%Y %H:%M:%S",
249
        handlers=[logging.StreamHandler(sys.stdout)],
Sylvain Gugger's avatar
Sylvain Gugger committed
250
    )
251

252
253
254
255
    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()

256
257
258
259
260
261
    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()
Sylvain Gugger's avatar
Sylvain Gugger committed
262
263
264

    # Log on each process the small summary:
    logger.warning(
265
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
266
        + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}"
Sylvain Gugger's avatar
Sylvain Gugger committed
267
    )
268
    logger.info(f"Training/evaluation parameters {training_args}")
Sylvain Gugger's avatar
Sylvain Gugger committed
269

270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
    # 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."
            )

Sylvain Gugger's avatar
Sylvain Gugger committed
285
286
287
288
289
290
291
292
293
294
295
296
297
298
    # 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 if no column called
    # 'text' is found. 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.
299
        raw_datasets = load_dataset(
300
301
302
            data_args.dataset_name,
            data_args.dataset_config_name,
            cache_dir=model_args.cache_dir,
303
            token=model_args.token,
304
            trust_remote_code=model_args.trust_remote_code,
305
        )
Sylvain Gugger's avatar
Sylvain Gugger committed
306
307
308
309
    else:
        data_files = {}
        if data_args.train_file is not None:
            data_files["train"] = data_args.train_file
310
311
            extension = data_args.train_file.split(".")[-1]

Sylvain Gugger's avatar
Sylvain Gugger committed
312
313
        if data_args.validation_file is not None:
            data_files["validation"] = data_args.validation_file
314
            extension = data_args.validation_file.split(".")[-1]
315
316
317
        if data_args.test_file is not None:
            data_files["test"] = data_args.test_file
            extension = data_args.test_file.split(".")[-1]
318
319
320
321
322
        raw_datasets = load_dataset(
            extension,
            data_files=data_files,
            field="data",
            cache_dir=model_args.cache_dir,
323
            token=model_args.token,
324
        )
Sylvain Gugger's avatar
Sylvain Gugger committed
325
    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
326
    # https://huggingface.co/docs/datasets/loading_datasets.
Sylvain Gugger's avatar
Sylvain Gugger committed
327
328
329
330
331
332
333
334
335

    # Load pretrained model and tokenizer
    #
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.
    config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
336
        revision=model_args.model_revision,
337
        token=model_args.token,
338
        trust_remote_code=model_args.trust_remote_code,
Sylvain Gugger's avatar
Sylvain Gugger committed
339
340
341
342
343
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        use_fast=True,
344
        revision=model_args.model_revision,
345
        token=model_args.token,
346
        trust_remote_code=model_args.trust_remote_code,
Sylvain Gugger's avatar
Sylvain Gugger committed
347
348
349
350
351
352
    )
    model = AutoModelForQuestionAnswering.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,
353
        revision=model_args.model_revision,
354
        token=model_args.token,
355
        trust_remote_code=model_args.trust_remote_code,
Sylvain Gugger's avatar
Sylvain Gugger committed
356
357
358
359
360
    )

    # Tokenizer check: this script requires a fast tokenizer.
    if not isinstance(tokenizer, PreTrainedTokenizerFast):
        raise ValueError(
Sylvain Gugger's avatar
Sylvain Gugger committed
361
362
363
            "This example script only works for models that have a fast tokenizer. Checkout the big table of models at"
            " https://huggingface.co/transformers/index.html#supported-frameworks to find the model types that meet"
            " this requirement"
Sylvain Gugger's avatar
Sylvain Gugger committed
364
365
366
        )

    # Preprocessing the datasets.
367
    # Preprocessing is slightly different for training and evaluation.
Sylvain Gugger's avatar
Sylvain Gugger committed
368
    if training_args.do_train:
369
        column_names = raw_datasets["train"].column_names
370
    elif training_args.do_eval:
371
        column_names = raw_datasets["validation"].column_names
372
    else:
373
        column_names = raw_datasets["test"].column_names
Sylvain Gugger's avatar
Sylvain Gugger committed
374
375
376
377
378
379
380
    question_column_name = "question" if "question" in column_names else column_names[0]
    context_column_name = "context" if "context" in column_names else column_names[1]
    answer_column_name = "answers" if "answers" in column_names else column_names[2]

    # Padding side determines if we do (question|context) or (context|question).
    pad_on_right = tokenizer.padding_side == "right"

381
    if data_args.max_seq_length > tokenizer.model_max_length:
382
        logger.warning(
383
            f"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the "
384
385
386
387
            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)

Sylvain Gugger's avatar
Sylvain Gugger committed
388
389
    # Training preprocessing
    def prepare_train_features(examples):
390
391
392
393
394
        # Some of the questions have lots of whitespace on the left, which is not useful and will make the
        # truncation of the context fail (the tokenized question will take a lots of space). So we remove that
        # left whitespace
        examples[question_column_name] = [q.lstrip() for q in examples[question_column_name]]

Sylvain Gugger's avatar
Sylvain Gugger committed
395
396
397
398
399
400
401
        # Tokenize our examples with truncation and maybe padding, but keep the overflows using a stride. This results
        # in one example possible giving several features when a context is long, each of those features having a
        # context that overlaps a bit the context of the previous feature.
        tokenized_examples = tokenizer(
            examples[question_column_name if pad_on_right else context_column_name],
            examples[context_column_name if pad_on_right else question_column_name],
            truncation="only_second" if pad_on_right else "only_first",
402
            max_length=max_seq_length,
Sylvain Gugger's avatar
Sylvain Gugger committed
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
            stride=data_args.doc_stride,
            return_overflowing_tokens=True,
            return_offsets_mapping=True,
            padding="max_length" if data_args.pad_to_max_length else False,
        )

        # Since one example might give us several features if it has a long context, we need a map from a feature to
        # its corresponding example. This key gives us just that.
        sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
        # The offset mappings will give us a map from token to character position in the original context. This will
        # help us compute the start_positions and end_positions.
        offset_mapping = tokenized_examples.pop("offset_mapping")

        # Let's label those examples!
        tokenized_examples["start_positions"] = []
        tokenized_examples["end_positions"] = []

        for i, offsets in enumerate(offset_mapping):
            # We will label impossible answers with the index of the CLS token.
            input_ids = tokenized_examples["input_ids"][i]
Matt's avatar
Matt committed
423
424
425
426
427
428
            if tokenizer.cls_token_id in input_ids:
                cls_index = input_ids.index(tokenizer.cls_token_id)
            elif tokenizer.bos_token_id in input_ids:
                cls_index = input_ids.index(tokenizer.bos_token_id)
            else:
                cls_index = 0
Sylvain Gugger's avatar
Sylvain Gugger committed
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471

            # Grab the sequence corresponding to that example (to know what is the context and what is the question).
            sequence_ids = tokenized_examples.sequence_ids(i)

            # One example can give several spans, this is the index of the example containing this span of text.
            sample_index = sample_mapping[i]
            answers = examples[answer_column_name][sample_index]
            # If no answers are given, set the cls_index as answer.
            if len(answers["answer_start"]) == 0:
                tokenized_examples["start_positions"].append(cls_index)
                tokenized_examples["end_positions"].append(cls_index)
            else:
                # Start/end character index of the answer in the text.
                start_char = answers["answer_start"][0]
                end_char = start_char + len(answers["text"][0])

                # Start token index of the current span in the text.
                token_start_index = 0
                while sequence_ids[token_start_index] != (1 if pad_on_right else 0):
                    token_start_index += 1

                # End token index of the current span in the text.
                token_end_index = len(input_ids) - 1
                while sequence_ids[token_end_index] != (1 if pad_on_right else 0):
                    token_end_index -= 1

                # Detect if the answer is out of the span (in which case this feature is labeled with the CLS index).
                if not (offsets[token_start_index][0] <= start_char and offsets[token_end_index][1] >= end_char):
                    tokenized_examples["start_positions"].append(cls_index)
                    tokenized_examples["end_positions"].append(cls_index)
                else:
                    # Otherwise move the token_start_index and token_end_index to the two ends of the answer.
                    # Note: we could go after the last offset if the answer is the last word (edge case).
                    while token_start_index < len(offsets) and offsets[token_start_index][0] <= start_char:
                        token_start_index += 1
                    tokenized_examples["start_positions"].append(token_start_index - 1)
                    while offsets[token_end_index][1] >= end_char:
                        token_end_index -= 1
                    tokenized_examples["end_positions"].append(token_end_index + 1)

        return tokenized_examples

    if training_args.do_train:
472
        if "train" not in raw_datasets:
473
            raise ValueError("--do_train requires a train dataset")
474
        train_dataset = raw_datasets["train"]
475
        if data_args.max_train_samples is not None:
Akul Agrawal's avatar
Akul Agrawal committed
476
            # We will select sample from whole data if argument is specified
477
478
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
            train_dataset = train_dataset.select(range(max_train_samples))
479
        # Create train feature from dataset
480
481
482
483
484
485
486
487
488
        with training_args.main_process_first(desc="train dataset map pre-processing"):
            train_dataset = train_dataset.map(
                prepare_train_features,
                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 train dataset",
            )
489
490
        if data_args.max_train_samples is not None:
            # Number of samples might increase during Feature Creation, We select only specified max samples
491
492
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
            train_dataset = train_dataset.select(range(max_train_samples))
Sylvain Gugger's avatar
Sylvain Gugger committed
493
494
495

    # Validation preprocessing
    def prepare_validation_features(examples):
496
497
498
499
500
        # Some of the questions have lots of whitespace on the left, which is not useful and will make the
        # truncation of the context fail (the tokenized question will take a lots of space). So we remove that
        # left whitespace
        examples[question_column_name] = [q.lstrip() for q in examples[question_column_name]]

Sylvain Gugger's avatar
Sylvain Gugger committed
501
502
503
504
505
506
507
        # Tokenize our examples with truncation and maybe padding, but keep the overflows using a stride. This results
        # in one example possible giving several features when a context is long, each of those features having a
        # context that overlaps a bit the context of the previous feature.
        tokenized_examples = tokenizer(
            examples[question_column_name if pad_on_right else context_column_name],
            examples[context_column_name if pad_on_right else question_column_name],
            truncation="only_second" if pad_on_right else "only_first",
508
            max_length=max_seq_length,
Sylvain Gugger's avatar
Sylvain Gugger committed
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
            stride=data_args.doc_stride,
            return_overflowing_tokens=True,
            return_offsets_mapping=True,
            padding="max_length" if data_args.pad_to_max_length else False,
        )

        # Since one example might give us several features if it has a long context, we need a map from a feature to
        # its corresponding example. This key gives us just that.
        sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")

        # For evaluation, we will need to convert our predictions to substrings of the context, so we keep the
        # corresponding example_id and we will store the offset mappings.
        tokenized_examples["example_id"] = []

        for i in range(len(tokenized_examples["input_ids"])):
            # Grab the sequence corresponding to that example (to know what is the context and what is the question).
            sequence_ids = tokenized_examples.sequence_ids(i)
            context_index = 1 if pad_on_right else 0

            # One example can give several spans, this is the index of the example containing this span of text.
            sample_index = sample_mapping[i]
            tokenized_examples["example_id"].append(examples["id"][sample_index])

            # Set to None the offset_mapping that are not part of the context so it's easy to determine if a token
            # position is part of the context or not.
            tokenized_examples["offset_mapping"][i] = [
                (o if sequence_ids[k] == context_index else None)
                for k, o in enumerate(tokenized_examples["offset_mapping"][i])
            ]

        return tokenized_examples

    if training_args.do_eval:
542
        if "validation" not in raw_datasets:
543
            raise ValueError("--do_eval requires a validation dataset")
544
        eval_examples = raw_datasets["validation"]
545
        if data_args.max_eval_samples is not None:
546
            # We will select sample from whole data
547
548
            max_eval_samples = min(len(eval_examples), data_args.max_eval_samples)
            eval_examples = eval_examples.select(range(max_eval_samples))
549
        # Validation Feature Creation
550
551
552
553
554
555
556
557
558
        with training_args.main_process_first(desc="validation dataset map pre-processing"):
            eval_dataset = eval_examples.map(
                prepare_validation_features,
                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 validation dataset",
            )
559
        if data_args.max_eval_samples is not None:
560
            # During Feature creation dataset samples might increase, we will select required samples again
561
562
            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
            eval_dataset = eval_dataset.select(range(max_eval_samples))
Sylvain Gugger's avatar
Sylvain Gugger committed
563

564
    if training_args.do_predict:
565
        if "test" not in raw_datasets:
566
            raise ValueError("--do_predict requires a test dataset")
567
        predict_examples = raw_datasets["test"]
568
        if data_args.max_predict_samples is not None:
569
            # We will select sample from whole data
570
571
            predict_examples = predict_examples.select(range(data_args.max_predict_samples))
        # Predict Feature Creation
572
573
574
575
576
577
578
579
580
        with training_args.main_process_first(desc="prediction dataset map pre-processing"):
            predict_dataset = predict_examples.map(
                prepare_validation_features,
                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 prediction dataset",
            )
581
        if data_args.max_predict_samples is not None:
582
            # During Feature creation dataset samples might increase, we will select required samples again
583
584
            max_predict_samples = min(len(predict_dataset), data_args.max_predict_samples)
            predict_dataset = predict_dataset.select(range(max_predict_samples))
585

Sylvain Gugger's avatar
Sylvain Gugger committed
586
587
588
    # Data collator
    # We have already padded to max length if the corresponding flag is True, otherwise we need to pad in the data
    # collator.
589
590
591
592
593
    data_collator = (
        default_data_collator
        if data_args.pad_to_max_length
        else DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8 if training_args.fp16 else None)
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
594
595

    # Post-processing:
596
    def post_processing_function(examples, features, predictions, stage="eval"):
Sylvain Gugger's avatar
Sylvain Gugger committed
597
598
599
600
601
602
603
604
605
606
        # Post-processing: we match the start logits and end logits to answers in the original context.
        predictions = postprocess_qa_predictions(
            examples=examples,
            features=features,
            predictions=predictions,
            version_2_with_negative=data_args.version_2_with_negative,
            n_best_size=data_args.n_best_size,
            max_answer_length=data_args.max_answer_length,
            null_score_diff_threshold=data_args.null_score_diff_threshold,
            output_dir=training_args.output_dir,
607
            log_level=log_level,
608
            prefix=stage,
Sylvain Gugger's avatar
Sylvain Gugger committed
609
610
611
612
        )
        # Format the result to the format the metric expects.
        if data_args.version_2_with_negative:
            formatted_predictions = [
613
                {"id": str(k), "prediction_text": v, "no_answer_probability": 0.0} for k, v in predictions.items()
Sylvain Gugger's avatar
Sylvain Gugger committed
614
615
            ]
        else:
616
            formatted_predictions = [{"id": str(k), "prediction_text": v} for k, v in predictions.items()]
617

618
        references = [{"id": str(ex["id"]), "answers": ex[answer_column_name]} for ex in examples]
Sylvain Gugger's avatar
Sylvain Gugger committed
619
620
        return EvalPrediction(predictions=formatted_predictions, label_ids=references)

621
622
623
624
625
626
627
628
    if data_args.version_2_with_negative:
        accepted_best_metrics = ("exact", "f1", "HasAns_exact", "HasAns_f1")
    else:
        accepted_best_metrics = ("exact_match", "f1")

    if training_args.load_best_model_at_end and training_args.metric_for_best_model not in accepted_best_metrics:
        warnings.warn(f"--metric_for_best_model should be set to one of {accepted_best_metrics}")

629
630
631
    metric = evaluate.load(
        "squad_v2" if data_args.version_2_with_negative else "squad", cache_dir=model_args.cache_dir
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
632
633
634
635
636
637
638
639
640

    def compute_metrics(p: EvalPrediction):
        return metric.compute(predictions=p.predictions, references=p.label_ids)

    # Initialize our Trainer
    trainer = QuestionAnsweringTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
641
        eval_dataset=eval_dataset if training_args.do_eval else None,
642
        eval_examples=eval_examples if training_args.do_eval else None,
Sylvain Gugger's avatar
Sylvain Gugger committed
643
644
645
646
647
648
649
650
        tokenizer=tokenizer,
        data_collator=data_collator,
        post_process_function=post_processing_function,
        compute_metrics=compute_metrics,
    )

    # Training
    if training_args.do_train:
651
652
653
654
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
655
656
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
Sylvain Gugger's avatar
Sylvain Gugger committed
657
658
        trainer.save_model()  # Saves the tokenizer too for easy upload

659
        metrics = train_result.metrics
660
661
662
663
664
        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))

665
666
667
        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()
668

Sylvain Gugger's avatar
Sylvain Gugger committed
669
670
671
    # Evaluation
    if training_args.do_eval:
        logger.info("*** Evaluate ***")
672
        metrics = trainer.evaluate()
Sylvain Gugger's avatar
Sylvain Gugger committed
673

674
675
        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))
Sylvain Gugger's avatar
Sylvain Gugger committed
676

677
678
        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)
Sylvain Gugger's avatar
Sylvain Gugger committed
679

680
681
682
    # Prediction
    if training_args.do_predict:
        logger.info("*** Predict ***")
683
        results = trainer.predict(predict_dataset, predict_examples)
684
685
        metrics = results.metrics

686
687
688
689
        max_predict_samples = (
            data_args.max_predict_samples if data_args.max_predict_samples is not None else len(predict_dataset)
        )
        metrics["predict_samples"] = min(max_predict_samples, len(predict_dataset))
690

691
692
        trainer.log_metrics("predict", metrics)
        trainer.save_metrics("predict", metrics)
693

694
695
696
697
698
699
700
701
    kwargs = {"finetuned_from": model_args.model_name_or_path, "tasks": "question-answering"}
    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
702

703
    if training_args.push_to_hub:
Sylvain Gugger's avatar
Sylvain Gugger committed
704
        trainer.push_to_hub(**kwargs)
705
706
    else:
        trainer.create_model_card(**kwargs)
Sylvain Gugger's avatar
Sylvain Gugger committed
707

Sylvain Gugger's avatar
Sylvain Gugger committed
708
709
710
711
712
713
714
715

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


if __name__ == "__main__":
    main()