run_squad.py 31.3 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,
                                  )
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
    'albert': (AlbertConfig, AlbertForQuestionAnswering, AlbertTokenizer),
thomwolf's avatar
thomwolf committed
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
def to_list(tensor):
    return tensor.detach().cpu().tolist()
thomwolf's avatar
thomwolf committed
77

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

83
    args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
thomwolf's avatar
thomwolf committed
84
85
86
87
    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:
88
        t_total = args.max_steps
thomwolf's avatar
thomwolf committed
89
90
        args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1
    else:
91
        t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs
thomwolf's avatar
thomwolf committed
92

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

thomwolf's avatar
thomwolf committed
102
103
104
105
106
    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
107
        
thomwolf's avatar
thomwolf committed
108
109
        model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)

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

thomwolf's avatar
thomwolf committed
114
115
116
117
118
119
    # 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
120
121
122
123
    # Train!
    logger.info("***** Running training *****")
    logger.info("  Num examples = %d", len(train_dataset))
    logger.info("  Num Epochs = %d", args.num_train_epochs)
124
125
126
    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
127
    logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
128
    logger.info("  Total optimization steps = %d", t_total)
thomwolf's avatar
thomwolf committed
129

Lysandre's avatar
Lysandre committed
130
    global_step = 1
thomwolf's avatar
thomwolf committed
131
    tr_loss, logging_loss = 0.0, 0.0
132
133
134
    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
135
    
136
137
138
139
    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
140
            batch = tuple(t.to(args.device) for t in batch)
LysandreJik's avatar
Cleanup  
LysandreJik committed
141
142
143
144

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

150
151
            if args.model_type in ['xlnet', 'xlm']:
                inputs.update({'cls_index': batch[5],
thomwolf's avatar
thomwolf committed
152
                               'p_mask':       batch[6]})
153
154
                if args.version_2_with_negative:
                    inputs.update({'is_impossible': batch[7]})
Peiqin Lin's avatar
typos  
Peiqin Lin committed
155
            outputs = model(**inputs)
156
            loss = outputs[0]  # model outputs are always tuple in transformers (see doc)
thomwolf's avatar
thomwolf committed
157

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

163
164
165
166
167
168
169
170
            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:
171
172
173
174
175
                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)

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

LysandreJik's avatar
Cleanup  
LysandreJik committed
181
                # Log metrics
182
183
184
185
186
187
188
189
190
                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
191
                # Save model checkpoint
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
                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
208
209
210
    if args.local_rank in [-1, 0]:
        tb_writer.close()

211
212
213
214
215
216
217
218
219
220
    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
221

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

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

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

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

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

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

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

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

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

278
            else:
LysandreJik's avatar
LysandreJik committed
279
280
281
282
283
                start_logits, end_logits = output
                result = SquadResult(
                    unique_id, start_logits, end_logits
                )

284
            all_results.append(result)
285

286
287
288
    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
289
    # Compute predictions
290
291
    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
292

293
294
295
296
    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
297

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

303
        predictions = compute_predictions_log_probs(examples, features, all_results, args.n_best_size,
304
                        args.max_answer_length, output_prediction_file,
LysandreJik's avatar
Cleanup  
LysandreJik committed
305
                        output_nbest_file, output_null_log_odds_file,
Lysandre's avatar
Lysandre committed
306
                        start_n_top, end_n_top,
307
                        args.version_2_with_negative, tokenizer, args.verbose_logging)
308
    else:
309
        predictions = compute_predictions_logits(examples, features, all_results, args.n_best_size,
310
311
                        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
312
                        args.version_2_with_negative, args.null_score_diff_threshold, tokenizer)
313

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

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

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

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

338
        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
339
340
341
342
343
344
345
346
347
348
349
350
            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()
351
352
353
354
355

            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
356

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

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

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

375
376
    if output_examples:
        return dataset, examples, features
thomwolf's avatar
thomwolf committed
377
378
    return dataset

379
380
381
382
383

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

    ## Required parameters
384
385
386
387
    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))
388
389
390
391
    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
392
    parser.add_argument("--data_dir", default=None, type=str,
393
394
395
396
397
398
399
400
                        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.")
401
402
403
404
405
406
407
    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
408
409
410
411
412
    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.")

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

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

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

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

    parser.add_argument('--threads', type=int, default=1, help='multiple threads for converting example to features')
485
486
    args = parser.parse_args()

thomwolf's avatar
thomwolf committed
487
488
489
    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))

490
    # Setup distant debugging if needed
491
492
493
494
495
496
497
    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
498
    # Setup CUDA, GPU & distributed training
499
500
    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
501
502
        args.n_gpu = torch.cuda.device_count()
    else:  # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
503
504
505
        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
506
507
        args.n_gpu = 1
    args.device = device
508

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

516
517
    # Set seed
    set_seed(args)
518

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

523
    args.model_type = args.model_type.lower()
524
    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
thomwolf's avatar
thomwolf committed
525
526
527
528
529
530
531
532
533
    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)
534
535

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

thomwolf's avatar
thomwolf committed
538
    model.to(args.device)
539

540
541
    logger.info("Training/evaluation parameters %s", args)

Simon Layton's avatar
Simon Layton committed
542
543
544
545
546
547
548
549
550
551
    # 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
552
    # Training
553
    if args.do_train:
554
        train_dataset = load_and_cache_examples(args, tokenizer, evaluate=False, output_examples=False)
555
        global_step, tr_loss = train(args, train_dataset, model, tokenizer)
556
        logger.info(" global_step = %s, average loss = %s", global_step, tr_loss)
557

558

thomwolf's avatar
thomwolf committed
559
    # Save the trained model and the tokenizer
Peng Qi's avatar
Peng Qi committed
560
    if args.do_train and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
561
562
563
564
565
566
567
568
569
570
        # 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)
571
572

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

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


thomwolf's avatar
thomwolf committed
581
    # Evaluation - we can ask to evaluate all the checkpoints (sub-directories) in a directory
582
583
    results = {}
    if args.do_eval and args.local_rank in [-1, 0]:
584
585
586
587
588
589
590
591
592
593

        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
594

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

597
        for checkpoint in checkpoints:
thomwolf's avatar
thomwolf committed
598
            # Reload the model
599
            global_step = checkpoint.split('-')[-1] if len(checkpoints) > 1 else ""
Lysandre's avatar
Lysandre committed
600
            model = model_class.from_pretrained(checkpoint, force_download=True)
601
            model.to(args.device)
thomwolf's avatar
thomwolf committed
602
603

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

606
607
            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
608

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

611
    return results
612
613
614
615


if __name__ == "__main__":
    main()