Add distilbert support (#126)

* add distilbert --------- Co-authored-by: Michael Wyatt <[email protected]>
microsoft · Aug 28, 2023 · 72b4898 · 72b4898
1 parent a011399
commit 72b4898
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -71,6 +71,7 @@ model family | size range | ~model count
 [gpt2](https://huggingface.co/models?other=gpt2) | 0.3B - 1.5B | 11,900
 [xlm-roberta](https://huggingface.co/models?other=xlm-roberta) | 0.1B - 0.3B | 4,100
 [roberta](https://huggingface.co/models?other=roberta) | 0.1B - 0.3B | 8,700
+[distilbert](https://huggingface.co/models?other=distilbert) | 0.1B - 0.3B | 4,700
 [bert](https://huggingface.co/models?other=bert) | 0.1B - 0.3B | 23,600
 
 <!--

diff --git a/mii/constants.py b/mii/constants.py
@@ -54,6 +54,7 @@ class ModelProvider(enum.Enum):
     'roberta': ModelProvider.HUGGING_FACE,
     'xlm-roberta': ModelProvider.HUGGING_FACE,
     'gpt2': ModelProvider.HUGGING_FACE,
+    'distilbert': ModelProvider.HUGGING_FACE,
     'bert': ModelProvider.HUGGING_FACE,
     'gpt_neo': ModelProvider.HUGGING_FACE,
     'gptj': ModelProvider.HUGGING_FACE,