run_squad.py 34.2 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
# Copyright (c) 2018, NVIDIA CORPORATION.  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.
16
""" Finetuning the library models for question-answering on SQuAD (DistilBERT, Bert, XLM, XLNet)."""
17
18

from __future__ import absolute_import, division, print_function
LysandreJik's avatar
LysandreJik committed
19
from transformers.data.processors.squad import SquadV1Processor, SquadV2Processor, SquadResult
20
from transformers.data.metrics.squad_metrics import compute_predictions_logits, compute_predictions_log_probs, squad_evaluate
21
22
23
24
25

import argparse
import logging
import os
import random
thomwolf's avatar
thomwolf committed
26
import glob
27
import timeit
28
29
import numpy as np
import torch
30
31
from torch.utils.data import (
    DataLoader, RandomSampler, SequentialSampler, TensorDataset)
32
33
from torch.utils.data.distributed import DistributedSampler

34
35
36
37
38
39
try:
    from torch.utils.tensorboard import SummaryWriter
except:
    from tensorboardX import SummaryWriter

from tqdm import tqdm, trange
40

41
from transformers import (WEIGHTS_NAME, BertConfig,
42
43
44
45
46
47
48
49
50
                          BertForQuestionAnswering, BertTokenizer,
                          XLMConfig, XLMForQuestionAnswering,
                          XLMTokenizer, XLNetConfig,
                          XLNetForQuestionAnswering,
                          XLNetTokenizer,
                          DistilBertConfig, DistilBertForQuestionAnswering, DistilBertTokenizer,
                          AlbertConfig, AlbertForQuestionAnswering, AlbertTokenizer,
                          XLMConfig, XLMForQuestionAnswering, XLMTokenizer,
                          )
thomwolf's avatar
thomwolf committed
51

Lysandre's avatar
Lysandre committed
52
from transformers import AdamW, get_linear_schedule_with_warmup, squad_convert_examples_to_features
53
54
55

logger = logging.getLogger(__name__)

56
ALL_MODELS = sum((tuple(conf.pretrained_config_archive_map.keys())
thomwolf's avatar
thomwolf committed
57
                  for conf in (BertConfig, XLNetConfig, XLMConfig)), ())
thomwolf's avatar
thomwolf committed
58
59

MODEL_CLASSES = {
thomwolf's avatar
thomwolf committed
60
61
62
    'bert': (BertConfig, BertForQuestionAnswering, BertTokenizer),
    'xlnet': (XLNetConfig, XLNetForQuestionAnswering, XLNetTokenizer),
    'xlm': (XLMConfig, XLMForQuestionAnswering, XLMTokenizer),
Lysandre's avatar
Lysandre committed
63
    'distilbert': (DistilBertConfig, DistilBertForQuestionAnswering, DistilBertTokenizer),
64
    'albert': (AlbertConfig, AlbertForQuestionAnswering, AlbertTokenizer),
thomwolf's avatar
thomwolf committed
65
66
}

67

thomwolf's avatar
thomwolf committed
68
69
70
71
72
73
74
def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)

75

76
77
def to_list(tensor):
    return tensor.detach().cpu().tolist()
thomwolf's avatar
thomwolf committed
78

79

80
def train(args, train_dataset, model, tokenizer):
thomwolf's avatar
thomwolf committed
81
82
83
84
    """ Train the model """
    if args.local_rank in [-1, 0]:
        tb_writer = SummaryWriter()

85
    args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
86
87
88
89
    train_sampler = RandomSampler(
        train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
    train_dataloader = DataLoader(
        train_dataset, sampler=train_sampler, batch_size=args.train_batch_size)
thomwolf's avatar
thomwolf committed
90
91

    if args.max_steps > 0:
92
        t_total = args.max_steps
93
94
        args.num_train_epochs = args.max_steps // (
            len(train_dataloader) // args.gradient_accumulation_steps) + 1
thomwolf's avatar
thomwolf committed
95
    else:
96
97
        t_total = len(
            train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs
thomwolf's avatar
thomwolf committed
98

99
    # Prepare optimizer and schedule (linear warmup and decay)
thomwolf's avatar
thomwolf committed
100
101
    no_decay = ['bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
102
103
104
105
        {'params': [p for n, p in model.named_parameters() if not any(
            nd in n for nd in no_decay)], 'weight_decay': args.weight_decay},
        {'params': [p for n, p in model.named_parameters() if any(
            nd in n for nd in no_decay)], 'weight_decay': 0.0}
LysandreJik's avatar
Cleanup  
LysandreJik committed
106
    ]
107
108
109
110
111
112
113
114
115
116
117
118
    optimizer = AdamW(optimizer_grouped_parameters,
                      lr=args.learning_rate, eps=args.adam_epsilon)
    scheduler = get_linear_schedule_with_warmup(
        optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total)

    # Check if saved optimizer or scheduler states exist
    if os.path.isfile(os.path.join(args.model_name_or_path, 'optimizer.pt')) and os.path.isfile(os.path.join(args.model_name_or_path, 'scheduler.pt')):
        # Load in optimizer and scheduler states
        optimizer.load_state_dict(torch.load(
            os.path.join(args.model_name_or_path, 'optimizer.pt')))
        scheduler.load_state_dict(torch.load(
            os.path.join(args.model_name_or_path, 'scheduler.pt')))
LysandreJik's avatar
Cleanup  
LysandreJik committed
119

thomwolf's avatar
thomwolf committed
120
121
122
123
    if args.fp16:
        try:
            from apex import amp
        except ImportError:
124
125
126
127
128
            raise ImportError(
                "Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")

        model, optimizer = amp.initialize(
            model, optimizer, opt_level=args.fp16_opt_level)
thomwolf's avatar
thomwolf committed
129

130
131
132
133
    # multi-gpu training (should be after apex fp16 initialization)
    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)

thomwolf's avatar
thomwolf committed
134
135
136
137
138
139
    # Distributed training (should be after apex fp16 initialization)
    if args.local_rank != -1:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],
                                                          output_device=args.local_rank,
                                                          find_unused_parameters=True)

thomwolf's avatar
thomwolf committed
140
141
142
143
    # Train!
    logger.info("***** Running training *****")
    logger.info("  Num examples = %d", len(train_dataset))
    logger.info("  Num Epochs = %d", args.num_train_epochs)
144
145
    logger.info("  Instantaneous batch size per GPU = %d",
                args.per_gpu_train_batch_size)
146
    logger.info("  Total train batch size (w. parallel, distributed & accumulation) = %d",
147
148
149
                args.train_batch_size * args.gradient_accumulation_steps * (torch.distributed.get_world_size() if args.local_rank != -1 else 1))
    logger.info("  Gradient Accumulation steps = %d",
                args.gradient_accumulation_steps)
150
    logger.info("  Total optimization steps = %d", t_total)
thomwolf's avatar
thomwolf committed
151

Lysandre's avatar
Lysandre committed
152
    global_step = 1
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
    epochs_trained = 0
    steps_trained_in_current_epoch = 0
    # Check if continuing training from a checkpoint
    if os.path.exists(args.model_name_or_path):
        # set global_step to gobal_step of last saved checkpoint from model path
        global_step = int(args.model_name_or_path.split('-')[-1].split('/')[0])
        epochs_trained = global_step // (len(train_dataloader) //
                                         args.gradient_accumulation_steps)
        steps_trained_in_current_epoch = global_step % (
            len(train_dataloader) // args.gradient_accumulation_steps)

        logger.info(
            "  Continuing training from checkpoint, will skip to saved global_step")
        logger.info("  Continuing training from epoch %d", epochs_trained)
        logger.info("  Continuing training from global step %d", global_step)
        logger.info("  Will skip the first %d steps in the first epoch",
                    steps_trained_in_current_epoch)

thomwolf's avatar
thomwolf committed
171
    tr_loss, logging_loss = 0.0, 0.0
172
    model.zero_grad()
173
174
175
176
177
    train_iterator = trange(epochs_trained, int(
        args.num_train_epochs), desc="Epoch", disable=args.local_rank not in [-1, 0])
    # Added here for reproductibility (even between python 2 and 3)
    set_seed(args)

178
    for _ in train_iterator:
179
180
        epoch_iterator = tqdm(train_dataloader, desc="Iteration",
                              disable=args.local_rank not in [-1, 0])
181
        for step, batch in enumerate(epoch_iterator):
182
183
184
185
186
187

            # Skip past any already trained steps if resuming training
            if steps_trained_in_current_epoch > 0:
                steps_trained_in_current_epoch -= 1
                continue

188
            model.train()
thomwolf's avatar
thomwolf committed
189
            batch = tuple(t.to(args.device) for t in batch)
LysandreJik's avatar
Cleanup  
LysandreJik committed
190
191
192
193
194
195
196
197

            inputs = {
                'input_ids':       batch[0],
                'attention_mask':  batch[1],
                'start_positions': batch[3],
                'end_positions':   batch[4]
            }

198
199
            if args.model_type != 'distilbert':
                inputs['token_type_ids'] = None if args.model_type == 'xlm' else batch[2]
LysandreJik's avatar
Cleanup  
LysandreJik committed
200

201
202
            if args.model_type in ['xlnet', 'xlm']:
                inputs.update({'cls_index': batch[5],
thomwolf's avatar
thomwolf committed
203
                               'p_mask':       batch[6]})
204
205
                if args.version_2_with_negative:
                    inputs.update({'is_impossible': batch[7]})
Peiqin Lin's avatar
typos  
Peiqin Lin committed
206
            outputs = model(**inputs)
207
208
            # model outputs are always tuple in transformers (see doc)
            loss = outputs[0]
thomwolf's avatar
thomwolf committed
209

210
            if args.n_gpu > 1:
211
                loss = loss.mean()  # mean() to average on multi-gpu parallel (not distributed) training
212
213
            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps
thomwolf's avatar
thomwolf committed
214

215
216
217
218
219
220
221
222
            if args.fp16:
                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                loss.backward()

            tr_loss += loss.item()
            if (step + 1) % args.gradient_accumulation_steps == 0:
223
                if args.fp16:
224
225
                    torch.nn.utils.clip_grad_norm_(
                        amp.master_params(optimizer), args.max_grad_norm)
226
                else:
227
228
                    torch.nn.utils.clip_grad_norm_(
                        model.parameters(), args.max_grad_norm)
229

230
                optimizer.step()
231
                scheduler.step()  # Update learning rate schedule
232
233
234
                model.zero_grad()
                global_step += 1

LysandreJik's avatar
Cleanup  
LysandreJik committed
235
                # Log metrics
236
                if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
237
238
                    # Only evaluate when single GPU otherwise metrics may not average well
                    if args.local_rank == -1 and args.evaluate_during_training:
239
240
                        results = evaluate(args, model, tokenizer)
                        for key, value in results.items():
241
242
243
244
245
246
                            tb_writer.add_scalar(
                                'eval_{}'.format(key), value, global_step)
                    tb_writer.add_scalar(
                        'lr', scheduler.get_lr()[0], global_step)
                    tb_writer.add_scalar(
                        'loss', (tr_loss - logging_loss)/args.logging_steps, global_step)
247
248
                    logging_loss = tr_loss

LysandreJik's avatar
Cleanup  
LysandreJik committed
249
                # Save model checkpoint
250
                if args.local_rank in [-1, 0] and args.save_steps > 0 and global_step % args.save_steps == 0:
251
252
                    output_dir = os.path.join(
                        args.output_dir, 'checkpoint-{}'.format(global_step))
253
254
                    if not os.path.exists(output_dir):
                        os.makedirs(output_dir)
255
256
257
                    # Take care of distributed/parallel training
                    model_to_save = model.module if hasattr(
                        model, 'module') else model
258
                    model_to_save.save_pretrained(output_dir)
259
260
261
262
                    tokenizer.save_pretrained(output_dir)

                    torch.save(args, os.path.join(
                        output_dir, 'training_args.bin'))
263
264
                    logger.info("Saving model checkpoint to %s", output_dir)

265
266
267
268
269
270
271
                    torch.save(optimizer.state_dict(), os.path.join(
                        output_dir, 'optimizer.pt'))
                    torch.save(scheduler.state_dict(), os.path.join(
                        output_dir, 'scheduler.pt'))
                    logger.info(
                        "Saving optimizer and scheduler states to %s", output_dir)

272
273
274
275
276
277
278
            if args.max_steps > 0 and global_step > args.max_steps:
                epoch_iterator.close()
                break
        if args.max_steps > 0 and global_step > args.max_steps:
            train_iterator.close()
            break

thomwolf's avatar
thomwolf committed
279
280
281
    if args.local_rank in [-1, 0]:
        tb_writer.close()

282
283
284
285
    return global_step, tr_loss / global_step


def evaluate(args, model, tokenizer, prefix=""):
286
287
    dataset, examples, features = load_and_cache_examples(
        args, tokenizer, evaluate=True, output_examples=True)
288
289
290
291
292

    if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:
        os.makedirs(args.output_dir)

    args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
LysandreJik's avatar
Cleanup  
LysandreJik committed
293

294
    # Note that DistributedSampler samples randomly
295
    eval_sampler = SequentialSampler(dataset)
296
297
    eval_dataloader = DataLoader(
        dataset, sampler=eval_sampler, batch_size=args.eval_batch_size)
298

ronakice's avatar
ronakice committed
299
    # multi-gpu evaluate
300
    if args.n_gpu > 1 and not isinstance(model, torch.nn.DataParallel):
ronakice's avatar
ronakice committed
301
302
        model = torch.nn.DataParallel(model)

303
304
305
306
    # Eval!
    logger.info("***** Running evaluation {} *****".format(prefix))
    logger.info("  Num examples = %d", len(dataset))
    logger.info("  Batch size = %d", args.eval_batch_size)
LysandreJik's avatar
Cleanup  
LysandreJik committed
307

308
    all_results = []
309
    start_time = timeit.default_timer()
LysandreJik's avatar
Cleanup  
LysandreJik committed
310

311
312
313
    for batch in tqdm(eval_dataloader, desc="Evaluating"):
        model.eval()
        batch = tuple(t.to(args.device) for t in batch)
LysandreJik's avatar
Cleanup  
LysandreJik committed
314

315
        with torch.no_grad():
LysandreJik's avatar
LysandreJik committed
316
317
318
319
            inputs = {
                'input_ids':      batch[0],
                'attention_mask': batch[1]
            }
320

321
            if args.model_type != 'distilbert':
322
323
                # XLM don't use segment_ids
                inputs['token_type_ids'] = None if args.model_type == 'xlm' else batch[2]
LysandreJik's avatar
Cleanup  
LysandreJik committed
324

325
            example_indices = batch[3]
326

LysandreJik's avatar
Cleanup  
LysandreJik committed
327
            # XLNet and XLM use more arguments for their predictions
328
            if args.model_type in ['xlnet', 'xlm']:
LysandreJik's avatar
Cleanup  
LysandreJik committed
329
330
                inputs.update({'cls_index': batch[4], 'p_mask': batch[5]})

331
332
333
334
335
            outputs = model(**inputs)

        for i, example_index in enumerate(example_indices):
            eval_feature = features[example_index.item()]
            unique_id = int(eval_feature.unique_id)
LysandreJik's avatar
LysandreJik committed
336

LysandreJik's avatar
LysandreJik committed
337
338
            output = [to_list(output[i]) for output in outputs]

LysandreJik's avatar
Cleanup  
LysandreJik committed
339
340
            # Some models (XLNet, XLM) use 5 arguments for their predictions, while the other "simpler"
            # models only use two.
LysandreJik's avatar
LysandreJik committed
341
342
343
344
            if len(output) >= 5:
                start_logits = output[0]
                start_top_index = output[1]
                end_logits = output[2]
LysandreJik's avatar
Cleanup  
LysandreJik committed
345
                end_top_index = output[3]
LysandreJik's avatar
LysandreJik committed
346
347
348
                cls_logits = output[4]

                result = SquadResult(
349
350
351
                    unique_id, start_logits, end_logits,
                    start_top_index=start_top_index,
                    end_top_index=end_top_index,
LysandreJik's avatar
LysandreJik committed
352
353
354
355
356
357
358
359
360
                    cls_logits=cls_logits
                )

            else:
                start_logits, end_logits = output
                result = SquadResult(
                    unique_id, start_logits, end_logits
                )

361
            all_results.append(result)
362

363
    evalTime = timeit.default_timer() - start_time
364
365
    logger.info("  Evaluation done in total %f secs (%f sec per example)",
                evalTime, evalTime / len(dataset))
366

thomwolf's avatar
thomwolf committed
367
    # Compute predictions
368
369
370
371
    output_prediction_file = os.path.join(
        args.output_dir, "predictions_{}.json".format(prefix))
    output_nbest_file = os.path.join(
        args.output_dir, "nbest_predictions_{}.json".format(prefix))
LysandreJik's avatar
Cleanup  
LysandreJik committed
372

373
    if args.version_2_with_negative:
374
375
        output_null_log_odds_file = os.path.join(
            args.output_dir, "null_odds_{}.json".format(prefix))
376
377
    else:
        output_null_log_odds_file = None
378

LysandreJik's avatar
Cleanup  
LysandreJik committed
379
    # XLNet and XLM use a more complex post-processing procedure
380
    if args.model_type in ['xlnet', 'xlm']:
381
382
383
384
        start_n_top = model.config.start_n_top if hasattr(
            model, "config") else model.module.config.start_n_top
        end_n_top = model.config.end_n_top if hasattr(
            model, "config") else model.module.config.end_n_top
Lysandre's avatar
Lysandre committed
385

386
        predictions = compute_predictions_log_probs(examples, features, all_results, args.n_best_size,
387
388
389
390
                                                    args.max_answer_length, output_prediction_file,
                                                    output_nbest_file, output_null_log_odds_file,
                                                    start_n_top, end_n_top,
                                                    args.version_2_with_negative, tokenizer, args.verbose_logging)
391
    else:
392
        predictions = compute_predictions_logits(examples, features, all_results, args.n_best_size,
393
394
395
                                                 args.max_answer_length, args.do_lower_case, output_prediction_file,
                                                 output_nbest_file, output_null_log_odds_file, args.verbose_logging,
                                                 args.version_2_with_negative, args.null_score_diff_threshold)
396

LysandreJik's avatar
Cleanup  
LysandreJik committed
397
    # Compute the F1 and exact scores.
LysandreJik's avatar
LysandreJik committed
398
    results = squad_evaluate(examples, predictions)
399
400
    return results

401

402
def load_and_cache_examples(args, tokenizer, evaluate=False, output_examples=False):
VictorSanh's avatar
VictorSanh committed
403
    if args.local_rank not in [-1, 0] and not evaluate:
404
405
        # Make sure only the first process in distributed training process the dataset, and the others will use the cache
        torch.distributed.barrier()
thomwolf's avatar
thomwolf committed
406

407
    # Load data features from cache or dataset file
LysandreJik's avatar
Cleanup  
LysandreJik committed
408
409
    input_dir = args.data_dir if args.data_dir else "."
    cached_features_file = os.path.join(input_dir, 'cached_{}_{}_{}'.format(
thomwolf's avatar
thomwolf committed
410
        'dev' if evaluate else 'train',
411
        list(filter(None, args.model_name_or_path.split('/'))).pop(),
LysandreJik's avatar
Cleanup  
LysandreJik committed
412
413
414
415
        str(args.max_seq_length))
    )

    # Init features and dataset from cache if it exists
416
    if os.path.exists(cached_features_file) and not args.overwrite_cache and not output_examples:
417
418
        logger.info("Loading features from cached file %s",
                    cached_features_file)
419
420
        features_and_dataset = torch.load(cached_features_file)
        features, dataset = features_and_dataset["features"], features_and_dataset["dataset"]
thomwolf's avatar
thomwolf committed
421
    else:
LysandreJik's avatar
Cleanup  
LysandreJik committed
422
        logger.info("Creating features from dataset file at %s", input_dir)
Lysandre's avatar
Lysandre committed
423

424
        if not args.data_dir and ((evaluate and not args.predict_file) or (not evaluate and not args.train_file)):
LysandreJik's avatar
Cleanup  
LysandreJik committed
425
426
427
            try:
                import tensorflow_datasets as tfds
            except ImportError:
428
429
                raise ImportError(
                    "If not data_dir is specified, tensorflow_datasets needs to be installed.")
LysandreJik's avatar
Cleanup  
LysandreJik committed
430
431

            if args.version_2_with_negative:
432
433
                logger.warn(
                    "tensorflow_datasets does not handle version 2 of SQuAD.")
LysandreJik's avatar
Cleanup  
LysandreJik committed
434
435

            tfds_examples = tfds.load("squad")
436
437
            examples = SquadV1Processor().get_examples_from_dataset(
                tfds_examples, evaluate=evaluate)
LysandreJik's avatar
Cleanup  
LysandreJik committed
438
439
        else:
            processor = SquadV2Processor() if args.version_2_with_negative else SquadV1Processor()
440
441
442
443
            if evaluate:
                examples = processor.get_dev_examples(args.data_dir, filename=args.predict_file)
            else:
                examples = processor.get_train_examples(args.data_dir, filename=args.train_file)
LysandreJik's avatar
LysandreJik committed
444

445
        features, dataset = squad_convert_examples_to_features(
Lysandre's avatar
Lysandre committed
446
447
448
449
450
451
            examples=examples,
            tokenizer=tokenizer,
            max_seq_length=args.max_seq_length,
            doc_stride=args.doc_stride,
            max_query_length=args.max_query_length,
            is_training=not evaluate,
452
            return_dataset='pt'
Lysandre's avatar
Lysandre committed
453
454
        )

thomwolf's avatar
thomwolf committed
455
        if args.local_rank in [-1, 0]:
456
457
458
459
            logger.info("Saving features into cached file %s",
                        cached_features_file)
            torch.save({"features": features, "dataset": dataset},
                       cached_features_file)
thomwolf's avatar
thomwolf committed
460

VictorSanh's avatar
VictorSanh committed
461
    if args.local_rank == 0 and not evaluate:
462
463
        # Make sure only the first process in distributed training process the dataset, and the others will use the cache
        torch.distributed.barrier()
thomwolf's avatar
thomwolf committed
464

465
466
    if output_examples:
        return dataset, examples, features
thomwolf's avatar
thomwolf committed
467
468
    return dataset

469
470
471
472

def main():
    parser = argparse.ArgumentParser()

473
    # Required parameters
474
475
476
477
    parser.add_argument("--model_type", default=None, type=str, required=True,
                        help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()))
    parser.add_argument("--model_name_or_path", default=None, type=str, required=True,
                        help="Path to pre-trained model or shortcut name selected in the list: " + ", ".join(ALL_MODELS))
478
479
480
    parser.add_argument("--output_dir", default=None, type=str, required=True,
                        help="The output directory where the model checkpoints and predictions will be written.")

481
    # Other parameters
LysandreJik's avatar
Cleanup  
LysandreJik committed
482
    parser.add_argument("--data_dir", default=None, type=str,
483
484
485
486
487
488
489
490
                        help="The input data dir. Should contain the .json files for the task." +
                             "If no data dir or train/predict files are specified, will run with tensorflow_datasets.")
    parser.add_argument("--train_file", default=None, type=str,
                        help="The input training file. If a data dir is specified, will look for the file there" +
                             "If no data dir or train/predict files are specified, will run with tensorflow_datasets.")
    parser.add_argument("--predict_file", default=None, type=str,
                        help="The input evaluation file. If a data dir is specified, will look for the file there" +
                             "If no data dir or train/predict files are specified, will run with tensorflow_datasets.")
491
492
493
494
495
496
497
    parser.add_argument("--config_name", default="", type=str,
                        help="Pretrained config name or path if not the same as model_name")
    parser.add_argument("--tokenizer_name", default="", type=str,
                        help="Pretrained tokenizer name or path if not the same as model_name")
    parser.add_argument("--cache_dir", default="", type=str,
                        help="Where do you want to store the pre-trained models downloaded from s3")

thomwolf's avatar
thomwolf committed
498
499
500
501
502
    parser.add_argument('--version_2_with_negative', action='store_true',
                        help='If true, the SQuAD examples contain some that do not have an answer.')
    parser.add_argument('--null_score_diff_threshold', type=float, default=0.0,
                        help="If null_score - best_non_null is greater than the threshold predict null.")

503
504
505
506
507
508
509
510
    parser.add_argument("--max_seq_length", default=384, type=int,
                        help="The maximum total input sequence length after WordPiece tokenization. Sequences "
                             "longer than this will be truncated, and sequences shorter than this will be padded.")
    parser.add_argument("--doc_stride", default=128, type=int,
                        help="When splitting up a long document into chunks, how much stride to take between chunks.")
    parser.add_argument("--max_query_length", default=64, type=int,
                        help="The maximum number of tokens for the question. Questions longer than this will "
                             "be truncated to this length.")
thomwolf's avatar
thomwolf committed
511
512
    parser.add_argument("--do_train", action='store_true',
                        help="Whether to run training.")
513
    parser.add_argument("--do_eval", action='store_true',
thomwolf's avatar
thomwolf committed
514
                        help="Whether to run eval on the dev set.")
515
516
    parser.add_argument("--evaluate_during_training", action='store_true',
                        help="Rul evaluation during training at each logging step.")
thomwolf's avatar
thomwolf committed
517
    parser.add_argument("--do_lower_case", action='store_true',
518
                        help="Set this flag if you are using an uncased model.")
thomwolf's avatar
thomwolf committed
519

520
521
522
523
    parser.add_argument("--per_gpu_train_batch_size", default=8, type=int,
                        help="Batch size per GPU/CPU for training.")
    parser.add_argument("--per_gpu_eval_batch_size", default=8, type=int,
                        help="Batch size per GPU/CPU for evaluation.")
thomwolf's avatar
thomwolf committed
524
525
526
527
    parser.add_argument("--learning_rate", default=5e-5, type=float,
                        help="The initial learning rate for Adam.")
    parser.add_argument('--gradient_accumulation_steps', type=int, default=1,
                        help="Number of updates steps to accumulate before performing a backward/update pass.")
528
    parser.add_argument("--weight_decay", default=0.0, type=float,
Kazutoshi Shinoda's avatar
Kazutoshi Shinoda committed
529
                        help="Weight decay if we apply some.")
530
531
532
533
    parser.add_argument("--adam_epsilon", default=1e-8, type=float,
                        help="Epsilon for Adam optimizer.")
    parser.add_argument("--max_grad_norm", default=1.0, type=float,
                        help="Max gradient norm.")
534
535
    parser.add_argument("--num_train_epochs", default=3.0, type=float,
                        help="Total number of training epochs to perform.")
536
537
538
539
    parser.add_argument("--max_steps", default=-1, type=int,
                        help="If > 0: set total number of training steps to perform. Override num_train_epochs.")
    parser.add_argument("--warmup_steps", default=0, type=int,
                        help="Linear warmup over warmup_steps.")
540
    parser.add_argument("--n_best_size", default=20, type=int,
thomwolf's avatar
thomwolf committed
541
                        help="The total number of n-best predictions to generate in the nbest_predictions.json output file.")
542
543
544
545
546
547
    parser.add_argument("--max_answer_length", default=30, type=int,
                        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.")
    parser.add_argument("--verbose_logging", action='store_true',
                        help="If true, all of the warnings related to data processing will be printed. "
                             "A number of warnings are expected for a normal SQuAD evaluation.")
thomwolf's avatar
thomwolf committed
548

549
550
551
552
553
554
    parser.add_argument('--logging_steps', type=int, default=50,
                        help="Log every X updates steps.")
    parser.add_argument('--save_steps', type=int, default=50,
                        help="Save checkpoint every X updates steps.")
    parser.add_argument("--eval_all_checkpoints", action='store_true',
                        help="Evaluate all checkpoints starting with the same prefix as model_name ending and ending with step number")
thomwolf's avatar
thomwolf committed
555
    parser.add_argument("--no_cuda", action='store_true',
556
                        help="Whether not to use CUDA when available")
557
558
559
560
    parser.add_argument('--overwrite_output_dir', action='store_true',
                        help="Overwrite the content of the output directory")
    parser.add_argument('--overwrite_cache', action='store_true',
                        help="Overwrite the cached training and evaluation sets")
thomwolf's avatar
thomwolf committed
561
    parser.add_argument('--seed', type=int, default=42,
562
                        help="random seed for initialization")
563

thomwolf's avatar
thomwolf committed
564
    parser.add_argument("--local_rank", type=int, default=-1,
565
                        help="local_rank for distributed training on gpus")
thomwolf's avatar
thomwolf committed
566
567
568
569
570
    parser.add_argument('--fp16', action='store_true',
                        help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit")
    parser.add_argument('--fp16_opt_level', type=str, default='O1',
                        help="For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
                             "See details at https://nvidia.github.io/apex/amp.html")
571
572
573
574
    parser.add_argument('--server_ip', type=str, default='',
                        help="Can be used for distant debugging.")
    parser.add_argument('--server_port', type=str, default='',
                        help="Can be used for distant debugging.")
575
576
    args = parser.parse_args()

thomwolf's avatar
thomwolf committed
577
    if os.path.exists(args.output_dir) and os.listdir(args.output_dir) and args.do_train and not args.overwrite_output_dir:
578
579
        raise ValueError(
            "Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome.".format(args.output_dir))
thomwolf's avatar
thomwolf committed
580

581
    # Setup distant debugging if needed
582
583
584
585
    if args.server_ip and args.server_port:
        # Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
        import ptvsd
        print("Waiting for debugger attach")
586
587
        ptvsd.enable_attach(
            address=(args.server_ip, args.server_port), redirect_output=True)
588
589
        ptvsd.wait_for_attach()

thomwolf's avatar
thomwolf committed
590
    # Setup CUDA, GPU & distributed training
591
    if args.local_rank == -1 or args.no_cuda:
592
593
        device = torch.device(
            "cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
thomwolf's avatar
thomwolf committed
594
595
        args.n_gpu = torch.cuda.device_count()
    else:  # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
596
597
598
        torch.cuda.set_device(args.local_rank)
        device = torch.device("cuda", args.local_rank)
        torch.distributed.init_process_group(backend='nccl')
thomwolf's avatar
thomwolf committed
599
600
        args.n_gpu = 1
    args.device = device
601

thomwolf's avatar
thomwolf committed
602
    # Setup logging
603
604
605
    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                        datefmt='%m/%d/%Y %H:%M:%S',
                        level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
thomwolf's avatar
thomwolf committed
606
    logger.warning("Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
607
                   args.local_rank, device, args.n_gpu, bool(args.local_rank != -1), args.fp16)
608

609
610
    # Set seed
    set_seed(args)
611

thomwolf's avatar
thomwolf committed
612
    # Load pretrained model and tokenizer
613
    if args.local_rank not in [-1, 0]:
614
615
        # Make sure only the first process in distributed training will download model & vocab
        torch.distributed.barrier()
616

617
    args.model_type = args.model_type.lower()
618
    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
thomwolf's avatar
thomwolf committed
619
620
621
622
623
624
    config = config_class.from_pretrained(args.config_name if args.config_name else args.model_name_or_path,
                                          cache_dir=args.cache_dir if args.cache_dir else None)
    tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name if args.tokenizer_name else args.model_name_or_path,
                                                do_lower_case=args.do_lower_case,
                                                cache_dir=args.cache_dir if args.cache_dir else None)
    model = model_class.from_pretrained(args.model_name_or_path,
625
626
                                        from_tf=bool(
                                            '.ckpt' in args.model_name_or_path),
thomwolf's avatar
thomwolf committed
627
628
                                        config=config,
                                        cache_dir=args.cache_dir if args.cache_dir else None)
629
630

    if args.local_rank == 0:
631
632
        # Make sure only the first process in distributed training will download model & vocab
        torch.distributed.barrier()
633

thomwolf's avatar
thomwolf committed
634
    model.to(args.device)
635

636
637
    logger.info("Training/evaluation parameters %s", args)

Simon Layton's avatar
Simon Layton committed
638
639
640
641
642
643
644
645
    # Before we do anything with models, we want to ensure that we get fp16 execution of torch.einsum if args.fp16 is set.
    # Otherwise it'll default to "promote" mode, and we'll get fp32 operations. Note that running `--fp16_opt_level="O2"` will
    # remove the need for this code, but it is still valid.
    if args.fp16:
        try:
            import apex
            apex.amp.register_half_function(torch, 'einsum')
        except ImportError:
646
647
            raise ImportError(
                "Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
Simon Layton's avatar
Simon Layton committed
648

thomwolf's avatar
thomwolf committed
649
    # Training
650
    if args.do_train:
651
652
        train_dataset = load_and_cache_examples(
            args, tokenizer, evaluate=False, output_examples=False)
653
        global_step, tr_loss = train(args, train_dataset, model, tokenizer)
654
655
        logger.info(" global_step = %s, average loss = %s",
                    global_step, tr_loss)
656

thomwolf's avatar
thomwolf committed
657
    # Save the trained model and the tokenizer
Peng Qi's avatar
Peng Qi committed
658
    if args.do_train and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
659
660
661
662
663
664
665
        # Create output directory if needed
        if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:
            os.makedirs(args.output_dir)

        logger.info("Saving model checkpoint to %s", args.output_dir)
        # Save a trained model, configuration and tokenizer using `save_pretrained()`.
        # They can then be reloaded using `from_pretrained()`
666
667
        # Take care of distributed/parallel training
        model_to_save = model.module if hasattr(model, 'module') else model
668
669
        model_to_save.save_pretrained(args.output_dir)
        tokenizer.save_pretrained(args.output_dir)
670
671

        # Good practice: save your training arguments together with the trained model
672
        torch.save(args, os.path.join(args.output_dir, 'training_args.bin'))
673

674
        # Load a trained model and vocabulary that you have fine-tuned
675
676
677
678
        model = model_class.from_pretrained(
            args.output_dir, force_download=True)
        tokenizer = tokenizer_class.from_pretrained(
            args.output_dir, do_lower_case=args.do_lower_case)
679
680
        model.to(args.device)

thomwolf's avatar
thomwolf committed
681
    # Evaluation - we can ask to evaluate all the checkpoints (sub-directories) in a directory
682
683
    results = {}
    if args.do_eval and args.local_rank in [-1, 0]:
684
685
686
687
688
689
690
691
692
        if args.do_train:
            logger.info("Loading checkpoints saved during training for evaluation")
            checkpoints = [args.output_dir]
            if args.eval_all_checkpoints:
                checkpoints = list(os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + '/**/' + WEIGHTS_NAME, recursive=True)))
                logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce model loading logs
        else:
            logger.info("Loading checkpoint %s for evaluation", args.model_name_or_path)
            checkpoints = [args.model_name_or_path]
thomwolf's avatar
thomwolf committed
693

694
        logger.info("Evaluate the following checkpoints: %s", checkpoints)
thomwolf's avatar
thomwolf committed
695

696
        for checkpoint in checkpoints:
thomwolf's avatar
thomwolf committed
697
            # Reload the model
698
699
700
701
            global_step = checkpoint.split(
                '-')[-1] if len(checkpoints) > 1 else ""
            model = model_class.from_pretrained(
                checkpoint, force_download=True)
702
            model.to(args.device)
thomwolf's avatar
thomwolf committed
703
704

            # Evaluate
705
            result = evaluate(args, model, tokenizer, prefix=global_step)
thomwolf's avatar
thomwolf committed
706

707
708
            result = dict((k + ('_{}'.format(global_step) if global_step else ''), v)
                          for k, v in result.items())
709
            results.update(result)
thomwolf's avatar
thomwolf committed
710

711
    logger.info("Results: {}".format(results))
thomwolf's avatar
thomwolf committed
712

713
    return results
714
715
716
717


if __name__ == "__main__":
    main()