Added DistilBERT to run_squad script

de8e14b6 · LysandreJik · 88368c2a · de8e14b6
Commit de8e14b6 authored Sep 11, 2019 by LysandreJik
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

examples/run_squad.py examples/run_squad.py +3 -1

No files found.
--- a/examples/run_squad.py
+++ b/examples/run_squad.py
@@ -37,7 +37,8 @@ from pytorch_transformers import (WEIGHTS_NAME, BertConfig,
                                  XLMConfig, XLMForQuestionAnswering,
                                  XLMTokenizer, XLNetConfig,
                                  XLNetForQuestionAnswering,
-                                  XLNetTokenizer)
+                                  XLNetTokenizer,
+                                  DistilBertConfig, DistilBertForQuestionAnswering, DistilBertTokenizer)

 from pytorch_transformers import AdamW, WarmupLinearSchedule

@@ -59,6 +60,7 @@ MODEL_CLASSES = {
    'bert': (BertConfig, BertForQuestionAnswering, BertTokenizer),
    'xlnet': (XLNetConfig, XLNetForQuestionAnswering, XLNetTokenizer),
    'xlm': (XLMConfig, XLMForQuestionAnswering, XLMTokenizer),
+    'distilbert': (DistilBertConfig, DistilBertForQuestionAnswering, DistilBertTokenizer)
 }

 def set_seed(args):