fix dataset names for distilbert (#1550)

Signed-off-by: Sheng Zha <[email protected]>
dmlc · Apr 29, 2021 · f822a73 · f822a73
1 parent 139b912
commit f822a73
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/scripts/bert/index.rst b/scripts/bert/index.rst
@@ -104,13 +104,13 @@ GluonNLP also supports the "`DistilBERT <https://arxiv.org/abs/1910.01108>`_" mo
 +-----------------------------------------+----------------------+
 |                                         | distilbert_6_768_12  |
 +=========================================+======================+
-| distil_book_corpus_wiki_en_uncased      | ✓                    |
+| distilbert_book_corpus_wiki_en_uncased  | ✓                    |
 +-----------------------------------------+----------------------+
 
 .. code-block:: python
 
     import gluonnlp as nlp; import mxnet as mx;
-    model, vocab = nlp.model.get_model('distilbert_6_768_12', dataset_name='distil_book_corpus_wiki_en_uncased');
+    model, vocab = nlp.model.get_model('distilbert_6_768_12', dataset_name='distilbert_book_corpus_wiki_en_uncased');
     tokenizer = nlp.data.BERTTokenizer(vocab, lower=True);
     transform = nlp.data.BERTSentenceTransform(tokenizer, max_seq_length=512, pair=False, pad=False);
     sample = transform(['Hello world!']);

diff --git a/src/gluonnlp/model/bert.py b/src/gluonnlp/model/bert.py
@@ -1254,7 +1254,7 @@ def bert_24_1024_16(dataset_name=None, vocab=None, pretrained=True, ctx=mx.cpu()
                           pretrained_allow_missing=pretrained_allow_missing,
                           hparam_allow_override=hparam_allow_override, **kwargs)
 
-def distilbert_6_768_12(dataset_name='distil_book_corpus_wiki_en_uncased', vocab=None,
+def distilbert_6_768_12(dataset_name='distilbert_book_corpus_wiki_en_uncased', vocab=None,
                         pretrained=True, ctx=mx.cpu(),
                         output_attention=False,
                         output_all_encodings=False,