run_squad.py 29.7 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
42
43
44
                                  BertForQuestionAnswering, BertTokenizer,
                                  XLMConfig, XLMForQuestionAnswering,
                                  XLMTokenizer, XLNetConfig,
                                  XLNetForQuestionAnswering,
45
                                  XLNetTokenizer,
Lysandre's avatar
Lysandre committed
46
47
                                  DistilBertConfig, DistilBertForQuestionAnswering, DistilBertTokenizer,
                                  AlbertConfig, AlbertForQuestionAnswering, AlbertTokenizer)
thomwolf's avatar
thomwolf committed
48

Lysandre's avatar
Lysandre committed
49
from transformers import AdamW, get_linear_schedule_with_warmup, squad_convert_examples_to_features
50
51
52

logger = logging.getLogger(__name__)

thomwolf's avatar
thomwolf committed
53
54
ALL_MODELS = sum((tuple(conf.pretrained_config_archive_map.keys()) \
                  for conf in (BertConfig, XLNetConfig, XLMConfig)), ())
thomwolf's avatar
thomwolf committed
55
56

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

thomwolf's avatar
thomwolf committed
64
65
66
67
68
69
70
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)

71
72
def to_list(tensor):
    return tensor.detach().cpu().tolist()
thomwolf's avatar
thomwolf committed
73

74
def train(args, train_dataset, model, tokenizer):
thomwolf's avatar
thomwolf committed
75
76
77
78
    """ Train the model """
    if args.local_rank in [-1, 0]:
        tb_writer = SummaryWriter()

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

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

thomwolf's avatar
thomwolf committed
98
99
100
101
102
    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
103
        
thomwolf's avatar
thomwolf committed
104
105
        model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)

106
107
108
109
    # multi-gpu training (should be after apex fp16 initialization)
    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)

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

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

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

145
146
            if args.model_type != 'distilbert':
                inputs['token_type_ids'] = None if args.model_type == 'xlm' else batch[2]
LysandreJik's avatar
Cleanup  
LysandreJik committed
147

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

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

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

159
160
161
162
163
164
165
166
            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:
167
168
169
170
171
                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)

172
                optimizer.step()
173
                scheduler.step()  # Update learning rate schedule
174
175
176
                model.zero_grad()
                global_step += 1

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

207
208
209
210
211
212
213
214
215
216
    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
217

218
    # Note that DistributedSampler samples randomly
219
    eval_sampler = SequentialSampler(dataset)
220
221
    eval_dataloader = DataLoader(dataset, sampler=eval_sampler, batch_size=args.eval_batch_size)

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

226
227
228
229
    # 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
230

231
    all_results = []
232
    start_time = timeit.default_timer()
LysandreJik's avatar
Cleanup  
LysandreJik committed
233

234
235
236
    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
237

238
        with torch.no_grad():
LysandreJik's avatar
LysandreJik committed
239
240
241
242
243
            inputs = {
                'input_ids':      batch[0],
                'attention_mask': batch[1]
            }
            
244
245
            if args.model_type != 'distilbert':
                inputs['token_type_ids'] = None if args.model_type == 'xlm' else batch[2]  # XLM don't use segment_ids
LysandreJik's avatar
Cleanup  
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']:
299
        predictions = compute_predictions_log_probs(examples, features, all_results, args.n_best_size,
300
                        args.max_answer_length, output_prediction_file,
LysandreJik's avatar
Cleanup  
LysandreJik committed
301
                        output_nbest_file, output_null_log_odds_file,
302
303
                        model.config.start_n_top, model.config.end_n_top,
                        args.version_2_with_negative, tokenizer, args.verbose_logging)
304
    else:
305
        predictions = compute_predictions_logits(examples, features, all_results, args.n_best_size,
306
307
308
                        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)
309

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

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

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

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

LysandreJik's avatar
Cleanup  
LysandreJik committed
334
335
336
337
338
339
340
341
342
343
344
345
346
347
        if not args.data_dir:
            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()
            examples = processor.get_dev_examples(args.data_dir) if evaluate else processor.get_train_examples(args.data_dir)
LysandreJik's avatar
LysandreJik committed
348

349
        features, dataset = squad_convert_examples_to_features( 
Lysandre's avatar
Lysandre committed
350
351
352
353
354
355
            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,
356
            return_dataset='pt'
Lysandre's avatar
Lysandre committed
357
358
        )

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

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

366
367
    if output_examples:
        return dataset, examples, features
thomwolf's avatar
thomwolf committed
368
369
    return dataset

370
371
372
373
374

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

    ## Required parameters
375
376
377
378
    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))
379
380
381
382
    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
383
384
    parser.add_argument("--data_dir", default=None, type=str,
                        help="The input data dir. Should contain the .json files for the task. If not specified, will run with tensorflow_datasets.")
385
386
387
388
389
390
391
    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
392
393
394
395
396
    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.")

397
398
399
400
401
402
403
404
    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
405
406
    parser.add_argument("--do_train", action='store_true',
                        help="Whether to run training.")
407
    parser.add_argument("--do_eval", action='store_true',
thomwolf's avatar
thomwolf committed
408
                        help="Whether to run eval on the dev set.")
409
410
    parser.add_argument("--evaluate_during_training", action='store_true',
                        help="Rul evaluation during training at each logging step.")
thomwolf's avatar
thomwolf committed
411
    parser.add_argument("--do_lower_case", action='store_true',
412
                        help="Set this flag if you are using an uncased model.")
thomwolf's avatar
thomwolf committed
413

414
415
416
417
    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
418
419
420
421
    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.")
422
    parser.add_argument("--weight_decay", default=0.0, type=float,
Kazutoshi Shinoda's avatar
Kazutoshi Shinoda committed
423
                        help="Weight decay if we apply some.")
424
425
426
427
    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.")
428
429
    parser.add_argument("--num_train_epochs", default=3.0, type=float,
                        help="Total number of training epochs to perform.")
430
431
432
433
    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.")
434
    parser.add_argument("--n_best_size", default=20, type=int,
thomwolf's avatar
thomwolf committed
435
                        help="The total number of n-best predictions to generate in the nbest_predictions.json output file.")
436
437
438
439
440
441
    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
442

443
444
445
446
447
448
    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
449
    parser.add_argument("--no_cuda", action='store_true',
450
                        help="Whether not to use CUDA when available")
451
452
453
454
    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
455
    parser.add_argument('--seed', type=int, default=42,
456
                        help="random seed for initialization")
457

thomwolf's avatar
thomwolf committed
458
    parser.add_argument("--local_rank", type=int, default=-1,
459
                        help="local_rank for distributed training on gpus")
thomwolf's avatar
thomwolf committed
460
461
462
463
464
    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")
465
466
467
468
    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()

LysandreJik's avatar
Cleanup  
LysandreJik committed
469
470
471
472
473
    args.predict_file = os.path.join(args.output_dir, 'predictions_{}_{}.txt'.format(
        list(filter(None, args.model_name_or_path.split('/'))).pop(),
        str(args.max_seq_length))
    )

thomwolf's avatar
thomwolf committed
474
475
476
    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))

477
    # Setup distant debugging if needed
478
479
480
481
482
483
484
    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
485
    # Setup CUDA, GPU & distributed training
486
487
    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
488
489
        args.n_gpu = torch.cuda.device_count()
    else:  # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
490
491
492
        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
493
494
        args.n_gpu = 1
    args.device = device
495

thomwolf's avatar
thomwolf committed
496
    # Setup logging
497
498
499
    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
500
    logger.warning("Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
501
                    args.local_rank, device, args.n_gpu, bool(args.local_rank != -1), args.fp16)
502

503
504
    # Set seed
    set_seed(args)
505

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

510
    args.model_type = args.model_type.lower()
511
    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
thomwolf's avatar
thomwolf committed
512
513
514
515
516
517
518
519
520
    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)
521
522

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

thomwolf's avatar
thomwolf committed
525
    model.to(args.device)
526

527
528
    logger.info("Training/evaluation parameters %s", args)

Simon Layton's avatar
Simon Layton committed
529
530
531
532
533
534
535
536
537
538
    # 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
539
    # Training
540
    if args.do_train:
541
        train_dataset = load_and_cache_examples(args, tokenizer, evaluate=False, output_examples=False)
542
        global_step, tr_loss = train(args, train_dataset, model, tokenizer)
543
        logger.info(" global_step = %s, average loss = %s", global_step, tr_loss)
544

545

thomwolf's avatar
thomwolf committed
546
    # Save the trained model and the tokenizer
Peng Qi's avatar
Peng Qi committed
547
    if args.do_train and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
548
549
550
551
552
553
554
555
556
557
        # 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)
558
559

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

562
        # Load a trained model and vocabulary that you have fine-tuned
Lysandre's avatar
Lysandre committed
563
        model = model_class.from_pretrained(args.output_dir, force_download=True)
Peng Qi's avatar
Peng Qi committed
564
        tokenizer = tokenizer_class.from_pretrained(args.output_dir, do_lower_case=args.do_lower_case)
565
566
567
        model.to(args.device)


thomwolf's avatar
thomwolf committed
568
    # Evaluation - we can ask to evaluate all the checkpoints (sub-directories) in a directory
569
570
571
572
573
    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)))
574
            logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce model loading logs
thomwolf's avatar
thomwolf committed
575

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

578
        for checkpoint in checkpoints:
thomwolf's avatar
thomwolf committed
579
            # Reload the model
580
            global_step = checkpoint.split('-')[-1] if len(checkpoints) > 1 else ""
Lysandre's avatar
Lysandre committed
581
            model = model_class.from_pretrained(checkpoint, force_download=True)
582
            model.to(args.device)
thomwolf's avatar
thomwolf committed
583
584

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

587
588
            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
589

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

592
    return results
593
594
595
596


if __name__ == "__main__":
    main()