Update run_glue.py

add DistilBert model shortcut into ALL_MODELS
2019-10-03 16:29:43 +08:00
parent c1689ac301
commit 7af0777910
1 changed files with 2 additions and 1 deletions
--- a/examples/run_glue.py
+++ b/examples/run_glue.py
@@ -53,7 +53,8 @@ from transformers import glue_convert_examples_to_features as convert_examples_t

 logger = logging.getLogger(__name__)

-ALL_MODELS = sum((tuple(conf.pretrained_config_archive_map.keys()) for conf in (BertConfig, XLNetConfig, XLMConfig, RobertaConfig)), ())
+ALL_MODELS = sum((tuple(conf.pretrained_config_archive_map.keys()) for conf in (BertConfig, XLNetConfig, XLMConfig, 
+                                                                                RobertaConfig, DistilBertConfig)), ())

 MODEL_CLASSES = {
    'bert': (BertConfig, BertForSequenceClassification, BertTokenizer),