Skip to content

GitLab

Commit 2a4ef098 authored by LysandreJik's avatar LysandreJik
Browse files

Add ALBERT and XLM to SQuAD script

parent 00c4e395
Hide whitespace changes
Inline Side-by-side
Showing with 5 additions and 2 deletions
examples/run_squad.py
View file @ 2a4ef098
...@@ -44,7 +44,9 @@ from transformers import (WEIGHTS_NAME, BertConfig, ...@@ -44,7 +44,9 @@ from transformers import (WEIGHTS_NAME, BertConfig,
XLNetForQuestionAnswering, XLNetForQuestionAnswering,
XLNetTokenizer, XLNetTokenizer,
DistilBertConfig, DistilBertForQuestionAnswering, DistilBertTokenizer, 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 from transformers import AdamW, get_linear_schedule_with_warmup, squad_convert_examples_to_features
...@@ -58,7 +60,8 @@ MODEL_CLASSES = { ...@@ -58,7 +60,8 @@ MODEL_CLASSES = {
'xlnet': (XLNetConfig, XLNetForQuestionAnswering, XLNetTokenizer), 'xlnet': (XLNetConfig, XLNetForQuestionAnswering, XLNetTokenizer),
'xlm': (XLMConfig, XLMForQuestionAnswering, XLMTokenizer), 'xlm': (XLMConfig, XLMForQuestionAnswering, XLMTokenizer),
'distilbert': (DistilBertConfig, DistilBertForQuestionAnswering, DistilBertTokenizer), 'distilbert': (DistilBertConfig, DistilBertForQuestionAnswering, DistilBertTokenizer),
'albert': (AlbertConfig, AlbertForQuestionAnswering, AlbertTokenizer) 'albert': (AlbertConfig, AlbertForQuestionAnswering, AlbertTokenizer),
'xlm': (XLMConfig, XLMForQuestionAnswering, XLMTokenizer)
} }
def set_seed(args): def set_seed(args):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment