SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

fix prediction on run-squad.py example

Browse Source
This commit is contained in:
joe dumoulin
2019-02-01 10:15:44 -08:00
parent 8f8bbd4a4c
commit aa90e0c36a
1 changed files with 7 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
examples/run_squad.py
View File

@@ -706,7 +706,7 @@ def main():
parser.add_argument("--num_train_epochs", default=3.0, type=float, parser.add_argument("--num_train_epochs", default=3.0, type=float,
help="Total number of training epochs to perform.") help="Total number of training epochs to perform.")
parser.add_argument("--warmup_proportion", default=0.1, type=float, parser.add_argument("--warmup_proportion", default=0.1, type=float,
help="Proportion of training to perform linear learning rate warmup for. E.g., 0.1 = 10%% " help="Proportion of training to perform linear learning rate warmup for. E.g., 0.1 = 10% "
"of training.") "of training.")
parser.add_argument("--n_best_size", default=20, type=int, parser.add_argument("--n_best_size", default=20, type=int,
help="The total number of n-best predictions to generate in the nbest_predictions.json " help="The total number of n-best predictions to generate in the nbest_predictions.json "
@@ -922,6 +922,9 @@ def main():
# Load a trained model that you have fine-tuned # Load a trained model that you have fine-tuned
model_state_dict = torch.load(output_model_file) model_state_dict = torch.load(output_model_file)
model = BertForQuestionAnswering.from_pretrained(args.bert_model, state_dict=model_state_dict) model = BertForQuestionAnswering.from_pretrained(args.bert_model, state_dict=model_state_dict)
else:
model = BertForQuestionAnswering.from_pretrained(args.bert_model)
model.to(device) model.to(device)
if args.do_predict and (args.local_rank == -1 or torch.distributed.get_rank() == 0): if args.do_predict and (args.local_rank == -1 or torch.distributed.get_rank() == 0):