run_squad.py 30.8 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
LysandreJik's avatar
Cleanup  
LysandreJik committed
30
from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler, TensorDataset)
31
32
from torch.utils.data.distributed import DistributedSampler

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

from tqdm import tqdm, trange
39

40
from transformers import (WEIGHTS_NAME, BertConfig,
thomwolf's avatar
thomwolf committed
41
                                  BertForQuestionAnswering, BertTokenizer,
erenup's avatar
erenup committed
42
                                  RobertaForQuestionAnswering, RobertaTokenizer, RobertaConfig,
thomwolf's avatar
thomwolf committed
43
44
45
                                  XLMConfig, XLMForQuestionAnswering,
                                  XLMTokenizer, XLNetConfig,
                                  XLNetForQuestionAnswering,
46
                                  XLNetTokenizer,
Lysandre's avatar
Lysandre committed
47
                                  DistilBertConfig, DistilBertForQuestionAnswering, DistilBertTokenizer,
48
49
50
                                  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__)

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

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

thomwolf's avatar
thomwolf committed
69
70
71
72
73
74
75
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)

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

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

84
    args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
thomwolf's avatar
thomwolf committed
85
86
87
88
    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)

    if args.max_steps > 0:
89
        t_total = args.max_steps
thomwolf's avatar
thomwolf committed
90
91
        args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1
    else:
92
        t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs
thomwolf's avatar
thomwolf committed
93

94
    # Prepare optimizer and schedule (linear warmup and decay)
thomwolf's avatar
thomwolf committed
95
96
    no_decay = ['bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
97
        {'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},
thomwolf's avatar
thomwolf committed
98
        {'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
99
    ]
100
    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
101
    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total)
LysandreJik's avatar
Cleanup  
LysandreJik committed
102

thomwolf's avatar
thomwolf committed
103
104
105
106
107
    if args.fp16:
        try:
            from apex import amp
        except ImportError:
            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
LysandreJik's avatar
Cleanup  
LysandreJik committed
108
        
thomwolf's avatar
thomwolf committed
109
110
        model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)

111
112
113
114
    # multi-gpu training (should be after apex fp16 initialization)
    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)

thomwolf's avatar
thomwolf committed
115
116
117
118
119
120
    # 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
121
122
123
124
    # Train!
    logger.info("***** Running training *****")
    logger.info("  Num examples = %d", len(train_dataset))
    logger.info("  Num Epochs = %d", args.num_train_epochs)
125
126
127
    logger.info("  Instantaneous batch size per GPU = %d", args.per_gpu_train_batch_size)
    logger.info("  Total train batch size (w. parallel, distributed & accumulation) = %d",
                   args.train_batch_size * args.gradient_accumulation_steps * (torch.distributed.get_world_size() if args.local_rank != -1 else 1))
thomwolf's avatar
thomwolf committed
128
    logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
129
    logger.info("  Total optimization steps = %d", t_total)
thomwolf's avatar
thomwolf committed
130

Lysandre's avatar
Lysandre committed
131
    global_step = 1
thomwolf's avatar
thomwolf committed
132
    tr_loss, logging_loss = 0.0, 0.0
133
134
135
    model.zero_grad()
    train_iterator = trange(int(args.num_train_epochs), desc="Epoch", disable=args.local_rank not in [-1, 0])
    set_seed(args)  # Added here for reproductibility (even between python 2 and 3)
LysandreJik's avatar
Cleanup  
LysandreJik committed
136
    
137
138
139
140
    for _ in train_iterator:
        epoch_iterator = tqdm(train_dataloader, desc="Iteration", disable=args.local_rank not in [-1, 0])
        for step, batch in enumerate(epoch_iterator):
            model.train()
thomwolf's avatar
thomwolf committed
141
            batch = tuple(t.to(args.device) for t in batch)
LysandreJik's avatar
Cleanup  
LysandreJik committed
142
143
144
145

            inputs = {
                'input_ids':       batch[0],
                'attention_mask':  batch[1],
erenup's avatar
erenup committed
146
                'token_type_ids': None if args.model_type in ['xlm', 'roberta', 'distilbert'] else batch[2],
LysandreJik's avatar
Cleanup  
LysandreJik committed
147
                'start_positions': batch[3],
erenup's avatar
erenup committed
148
                'end_positions':   batch[4],
LysandreJik's avatar
Cleanup  
LysandreJik committed
149
150
            }

151
            if args.model_type in ['xlnet', 'xlm']:
LysandreJik's avatar
Cleanup  
LysandreJik committed
152
153
                inputs.update({'cls_index': batch[5], 'p_mask': batch[6]})

Peiqin Lin's avatar
typos  
Peiqin Lin committed
154
            outputs = model(**inputs)
155
            loss = outputs[0]  # model outputs are always tuple in transformers (see doc)
thomwolf's avatar
thomwolf committed
156

157
            if args.n_gpu > 1:
thomwolf's avatar
thomwolf committed
158
                loss = loss.mean() # mean() to average on multi-gpu parallel (not distributed) training
159
160
            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps
thomwolf's avatar
thomwolf committed
161

162
163
164
165
166
167
168
169
            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:
170
171
172
173
174
                if args.fp16:
                    torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
                else:
                    torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)

175
                optimizer.step()
176
                scheduler.step()  # Update learning rate schedule
177
178
179
                model.zero_grad()
                global_step += 1

LysandreJik's avatar
Cleanup  
LysandreJik committed
180
                # Log metrics
181
182
183
184
185
186
187
188
189
                if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
                    if args.local_rank == -1 and args.evaluate_during_training:  # Only evaluate when single GPU otherwise metrics may not average well
                        results = evaluate(args, model, tokenizer)
                        for key, value in results.items():
                            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)
                    logging_loss = tr_loss

LysandreJik's avatar
Cleanup  
LysandreJik committed
190
                # Save model checkpoint
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
                if args.local_rank in [-1, 0] and args.save_steps > 0 and global_step % args.save_steps == 0:
                    output_dir = os.path.join(args.output_dir, 'checkpoint-{}'.format(global_step))
                    if not os.path.exists(output_dir):
                        os.makedirs(output_dir)
                    model_to_save = model.module if hasattr(model, 'module') else model  # Take care of distributed/parallel training
                    model_to_save.save_pretrained(output_dir)
                    torch.save(args, os.path.join(output_dir, 'training_args.bin'))
                    logger.info("Saving model checkpoint to %s", output_dir)

            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
207
208
209
    if args.local_rank in [-1, 0]:
        tb_writer.close()

210
211
212
213
214
215
216
217
218
219
    return global_step, tr_loss / global_step


def evaluate(args, model, tokenizer, prefix=""):
    dataset, examples, features = load_and_cache_examples(args, tokenizer, evaluate=True, output_examples=True)

    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
220

221
    # Note that DistributedSampler samples randomly
222
    eval_sampler = SequentialSampler(dataset)
223
224
    eval_dataloader = DataLoader(dataset, sampler=eval_sampler, batch_size=args.eval_batch_size)

ronakice's avatar
ronakice committed
225
226
227
228
    # multi-gpu evaluate
    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)

229
230
231
232
    # 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
233

234
    all_results = []
235
    start_time = timeit.default_timer()
LysandreJik's avatar
Cleanup  
LysandreJik committed
236

237
238
239
    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
240

241
        with torch.no_grad():
LysandreJik's avatar
LysandreJik committed
242
243
            inputs = {
                'input_ids':      batch[0],
erenup's avatar
erenup committed
244
245
                'attention_mask': batch[1],
                'token_type_ids': None if args.model_type in ['xlm', 'roberta', 'distilbert'] else batch[2],
LysandreJik's avatar
LysandreJik committed
246
            }
247
            example_indices = batch[3]
LysandreJik's avatar
Cleanup  
LysandreJik committed
248
249
            
            # XLNet and XLM use more arguments for their predictions
250
            if args.model_type in ['xlnet', 'xlm']:
LysandreJik's avatar
Cleanup  
LysandreJik committed
251
252
                inputs.update({'cls_index': batch[4], 'p_mask': batch[5]})

253
254
255
256
257
            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
258

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

LysandreJik's avatar
Cleanup  
LysandreJik committed
261
262
            # Some models (XLNet, XLM) use 5 arguments for their predictions, while the other "simpler"
            # models only use two.
LysandreJik's avatar
LysandreJik committed
263
264
265
266
            if len(output) >= 5:
                start_logits = output[0]
                start_top_index = output[1]
                end_logits = output[2]
LysandreJik's avatar
Cleanup  
LysandreJik committed
267
                end_top_index = output[3]
LysandreJik's avatar
LysandreJik committed
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
                cls_logits = output[4]

                result = SquadResult(
                    unique_id, start_logits, end_logits, 
                    start_top_index=start_top_index, 
                    end_top_index=end_top_index, 
                    cls_logits=cls_logits
                )

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

283
            all_results.append(result)
284

285
286
287
    evalTime = timeit.default_timer() - start_time
    logger.info("  Evaluation done in total %f secs (%f sec per example)", evalTime, evalTime / len(dataset))

thomwolf's avatar
thomwolf committed
288
    # Compute predictions
289
290
    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
291

292
293
294
295
    if args.version_2_with_negative:
        output_null_log_odds_file = os.path.join(args.output_dir, "null_odds_{}.json".format(prefix))
    else:
        output_null_log_odds_file = None
296

LysandreJik's avatar
Cleanup  
LysandreJik committed
297
    # XLNet and XLM use a more complex post-processing procedure
298
    if args.model_type in ['xlnet', 'xlm']:
Lysandre's avatar
Lysandre committed
299
300
301
        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

302
        predictions = compute_predictions_log_probs(examples, features, all_results, args.n_best_size,
303
                        args.max_answer_length, output_prediction_file,
LysandreJik's avatar
Cleanup  
LysandreJik committed
304
                        output_nbest_file, output_null_log_odds_file,
Lysandre's avatar
Lysandre committed
305
                        start_n_top, end_n_top,
306
                        args.version_2_with_negative, tokenizer, args.verbose_logging)
307
    else:
308
        predictions = compute_predictions_logits(examples, features, all_results, args.n_best_size,
309
310
                        args.max_answer_length, args.do_lower_case, output_prediction_file,
                        output_nbest_file, output_null_log_odds_file, args.verbose_logging,
erenup's avatar
erenup committed
311
                        args.version_2_with_negative, args.null_score_diff_threshold, tokenizer)
312

LysandreJik's avatar
Cleanup  
LysandreJik committed
313
    # Compute the F1 and exact scores.
LysandreJik's avatar
LysandreJik committed
314
    results = squad_evaluate(examples, predictions)
315
316
317
    return results

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

321
    # Load data features from cache or dataset file
LysandreJik's avatar
Cleanup  
LysandreJik committed
322
323
    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
324
        'dev' if evaluate else 'train',
325
        list(filter(None, args.model_name_or_path.split('/'))).pop(),
LysandreJik's avatar
Cleanup  
LysandreJik committed
326
327
328
329
        str(args.max_seq_length))
    )

    # Init features and dataset from cache if it exists
330
    if os.path.exists(cached_features_file) and not args.overwrite_cache and not output_examples:
thomwolf's avatar
thomwolf committed
331
        logger.info("Loading features from cached file %s", cached_features_file)
332
333
        features_and_dataset = torch.load(cached_features_file)
        features, dataset = features_and_dataset["features"], features_and_dataset["dataset"]
thomwolf's avatar
thomwolf committed
334
    else:
LysandreJik's avatar
Cleanup  
LysandreJik committed
335
        logger.info("Creating features from dataset file at %s", input_dir)
Lysandre's avatar
Lysandre committed
336

337
        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
338
339
340
341
342
343
344
345
346
347
348
349
            try:
                import tensorflow_datasets as tfds
            except ImportError:
                raise ImportError("If not data_dir is specified, tensorflow_datasets needs to be installed.")

            if args.version_2_with_negative:
                logger.warn("tensorflow_datasets does not handle version 2 of SQuAD.")

            tfds_examples = tfds.load("squad")
            examples = SquadV1Processor().get_examples_from_dataset(tfds_examples, evaluate=evaluate)
        else:
            processor = SquadV2Processor() if args.version_2_with_negative else SquadV1Processor()
350
351
352
353
354

            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
355

356
        features, dataset = squad_convert_examples_to_features( 
Lysandre's avatar
Lysandre committed
357
358
359
360
361
362
            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,
363
            return_dataset='pt'
Lysandre's avatar
Lysandre committed
364
365
        )

thomwolf's avatar
thomwolf committed
366
367
        if args.local_rank in [-1, 0]:
            logger.info("Saving features into cached file %s", cached_features_file)
368
            torch.save({"features": features, "dataset": dataset}, cached_features_file)
thomwolf's avatar
thomwolf committed
369

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

373
374
    if output_examples:
        return dataset, examples, features
thomwolf's avatar
thomwolf committed
375
376
    return dataset

377
378
379
380
381

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

    ## Required parameters
382
383
384
385
    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))
386
387
388
389
    parser.add_argument("--output_dir", default=None, type=str, required=True,
                        help="The output directory where the model checkpoints and predictions will be written.")

    ## Other parameters
LysandreJik's avatar
Cleanup  
LysandreJik committed
390
    parser.add_argument("--data_dir", default=None, type=str,
391
392
393
394
395
396
397
398
                        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.")
399
400
401
402
403
404
405
    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
406
407
408
409
410
    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.")

411
412
413
414
415
416
417
418
    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
419
420
    parser.add_argument("--do_train", action='store_true',
                        help="Whether to run training.")
421
    parser.add_argument("--do_eval", action='store_true',
thomwolf's avatar
thomwolf committed
422
                        help="Whether to run eval on the dev set.")
423
424
    parser.add_argument("--evaluate_during_training", action='store_true',
                        help="Rul evaluation during training at each logging step.")
thomwolf's avatar
thomwolf committed
425
    parser.add_argument("--do_lower_case", action='store_true',
426
                        help="Set this flag if you are using an uncased model.")
thomwolf's avatar
thomwolf committed
427

428
429
430
431
    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
432
433
434
435
    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.")
436
    parser.add_argument("--weight_decay", default=0.0, type=float,
Kazutoshi Shinoda's avatar
Kazutoshi Shinoda committed
437
                        help="Weight decay if we apply some.")
438
439
440
441
    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.")
442
443
    parser.add_argument("--num_train_epochs", default=3.0, type=float,
                        help="Total number of training epochs to perform.")
444
445
446
447
    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.")
448
    parser.add_argument("--n_best_size", default=20, type=int,
thomwolf's avatar
thomwolf committed
449
                        help="The total number of n-best predictions to generate in the nbest_predictions.json output file.")
450
451
452
453
454
455
    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
456

457
458
459
460
461
462
    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
463
    parser.add_argument("--no_cuda", action='store_true',
464
                        help="Whether not to use CUDA when available")
465
466
467
468
    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
469
    parser.add_argument('--seed', type=int, default=42,
470
                        help="random seed for initialization")
471

thomwolf's avatar
thomwolf committed
472
    parser.add_argument("--local_rank", type=int, default=-1,
473
                        help="local_rank for distributed training on gpus")
thomwolf's avatar
thomwolf committed
474
475
476
477
478
    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")
479
480
481
482
    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.")
    args = parser.parse_args()

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

486
    # Setup distant debugging if needed
487
488
489
490
491
492
493
    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")
        ptvsd.enable_attach(address=(args.server_ip, args.server_port), redirect_output=True)
        ptvsd.wait_for_attach()

thomwolf's avatar
thomwolf committed
494
    # Setup CUDA, GPU & distributed training
495
496
    if args.local_rank == -1 or args.no_cuda:
        device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
thomwolf's avatar
thomwolf committed
497
498
        args.n_gpu = torch.cuda.device_count()
    else:  # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
499
500
501
        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
502
503
        args.n_gpu = 1
    args.device = device
504

thomwolf's avatar
thomwolf committed
505
    # Setup logging
506
507
508
    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
509
    logger.warning("Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
510
                    args.local_rank, device, args.n_gpu, bool(args.local_rank != -1), args.fp16)
511

512
513
    # Set seed
    set_seed(args)
514

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

519
    args.model_type = args.model_type.lower()
520
    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
thomwolf's avatar
thomwolf committed
521
522
523
524
525
526
527
528
529
    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,
                                        from_tf=bool('.ckpt' in args.model_name_or_path),
                                        config=config,
                                        cache_dir=args.cache_dir if args.cache_dir else None)
530
531

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

thomwolf's avatar
thomwolf committed
534
    model.to(args.device)
535

536
537
    logger.info("Training/evaluation parameters %s", args)

Simon Layton's avatar
Simon Layton committed
538
539
540
541
542
543
544
545
546
547
    # 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:
            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")

thomwolf's avatar
thomwolf committed
548
    # Training
549
    if args.do_train:
550
        train_dataset = load_and_cache_examples(args, tokenizer, evaluate=False, output_examples=False)
551
        global_step, tr_loss = train(args, train_dataset, model, tokenizer)
552
        logger.info(" global_step = %s, average loss = %s", global_step, tr_loss)
553

554

thomwolf's avatar
thomwolf committed
555
    # Save the trained model and the tokenizer
Peng Qi's avatar
Peng Qi committed
556
    if args.do_train and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
557
558
559
560
561
562
563
564
565
566
        # 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()`
        model_to_save = model.module if hasattr(model, 'module') else model  # Take care of distributed/parallel training
        model_to_save.save_pretrained(args.output_dir)
        tokenizer.save_pretrained(args.output_dir)
567
568

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

571
        # Load a trained model and vocabulary that you have fine-tuned
Lysandre's avatar
Lysandre committed
572
        model = model_class.from_pretrained(args.output_dir, force_download=True)
Peng Qi's avatar
Peng Qi committed
573
        tokenizer = tokenizer_class.from_pretrained(args.output_dir, do_lower_case=args.do_lower_case)
574
575
576
        model.to(args.device)


thomwolf's avatar
thomwolf committed
577
    # Evaluation - we can ask to evaluate all the checkpoints (sub-directories) in a directory
578
579
580
581
582
    results = {}
    if args.do_eval and args.local_rank in [-1, 0]:
        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)))
583
            logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce model loading logs
thomwolf's avatar
thomwolf committed
584

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

587
        for checkpoint in checkpoints:
thomwolf's avatar
thomwolf committed
588
            # Reload the model
589
            global_step = checkpoint.split('-')[-1] if len(checkpoints) > 1 else ""
Lysandre's avatar
Lysandre committed
590
            model = model_class.from_pretrained(checkpoint, force_download=True)
591
            model.to(args.device)
thomwolf's avatar
thomwolf committed
592
593

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

596
597
            result = dict((k + ('_{}'.format(global_step) if global_step else ''), v) for k, v in result.items())
            results.update(result)
thomwolf's avatar
thomwolf committed
598

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

601
    return results
602
603
604
605


if __name__ == "__main__":
    main()