From 2a4ef098d65939d436e2a5efbb518fb807b6b1b6 Mon Sep 17 00:00:00 2001 From: LysandreJik Date: Mon, 9 Dec 2019 10:46:47 -0500 Subject: [PATCH] Add ALBERT and XLM to SQuAD script --- examples/run_squad.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/examples/run_squad.py b/examples/run_squad.py index a8ac1d1b05..2df29014ef 100644 --- a/examples/run_squad.py +++ b/examples/run_squad.py @@ -44,7 +44,9 @@ from transformers import (WEIGHTS_NAME, BertConfig, XLNetForQuestionAnswering, XLNetTokenizer, DistilBertConfig, DistilBertForQuestionAnswering, DistilBertTokenizer, - AlbertConfig, AlbertForQuestionAnswering, AlbertTokenizer) + AlbertConfig, AlbertForQuestionAnswering, AlbertTokenizer, + XLMConfig, XLMForQuestionAnswering, XLMTokenizer, + ) from transformers import AdamW, get_linear_schedule_with_warmup, squad_convert_examples_to_features @@ -58,7 +60,8 @@ MODEL_CLASSES = { 'xlnet': (XLNetConfig, XLNetForQuestionAnswering, XLNetTokenizer), 'xlm': (XLMConfig, XLMForQuestionAnswering, XLMTokenizer), 'distilbert': (DistilBertConfig, DistilBertForQuestionAnswering, DistilBertTokenizer), - 'albert': (AlbertConfig, AlbertForQuestionAnswering, AlbertTokenizer) + 'albert': (AlbertConfig, AlbertForQuestionAnswering, AlbertTokenizer), + 'xlm': (XLMConfig, XLMForQuestionAnswering, XLMTokenizer) } def set_seed(args):