include roberta in run_squad_w_distillation - cc @graviraja

852e032c · VictorSanh · b5509abb · 852e032c
Commit 852e032c authored Mar 01, 2020 by VictorSanh
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 0 deletions

examples/distillation/run_squad_w_distillation.py examples/distillation/run_squad_w_distillation.py +4 -0

No files found.
--- 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),
 }