include roberta in run_squad_w_distillation - cc @graviraja

2020-03-01 01:56:50 +00:00
parent b5509abb36
commit 852e032ca6
1 changed files with 4 additions and 0 deletions
--- a/examples/distillation/run_squad_w_distillation.py
+++ b/examples/distillation/run_squad_w_distillation.py
@@ -45,6 +45,9 @@ from transformers import (
    XLNetConfig,
    XLNetForQuestionAnswering,
    XLNetTokenizer,
+    RobertaConfig,
+    RobertaForQuestionAnswering,
+    RobertaTokenizer,
    get_linear_schedule_with_warmup,
    squad_convert_examples_to_features,
 )
@@ -73,6 +76,7 @@ MODEL_CLASSES = {
    "xlnet": (XLNetConfig, XLNetForQuestionAnswering, XLNetTokenizer),
    "xlm": (XLMConfig, XLMForQuestionAnswering, XLMTokenizer),
    "distilbert": (DistilBertConfig, DistilBertForQuestionAnswering, DistilBertTokenizer),
+    "roberta": (RobertaConfig, RobertaForQuestionAnswering, RobertaTokenizer),
 }