Add DistilBERT model (#172)

Co-authored-by: Jonatan Kłosko <jonatanklosko@gmail.com>
elixir-nx · Mar 10, 2023 · 2899800 · 2899800
1 parent fd9a8b1
commit 2899800
Show file tree

Hide file tree

Showing 5 changed files with 721 additions and 0 deletions.
diff --git a/lib/bumblebee.ex b/lib/bumblebee.ex
@@ -75,6 +75,12 @@ defmodule Bumblebee do
       {Bumblebee.Vision.Deit, :for_image_classification_with_teacher},
     "DeiTForMaskedImageModeling" => {Bumblebee.Vision.Deit, :for_masked_image_modeling},
     "DeiTModel" => {Bumblebee.Vision.Deit, :base},
+    "DistilBertModel" => {Bumblebee.Text.Distilbert, :base},
+    "DistilBertForMaskedLM" => {Bumblebee.Text.Distilbert, :for_masked_language_modeling},
+    "DistilBertForSequenceClassification" =>
+      {Bumblebee.Text.Distilbert, :for_sequence_classification},
+    "DistilBertForQuestionAnswering" => {Bumblebee.Text.Distilbert, :for_question_answering},
+    "DistilBertForTokenClassification" => {Bumblebee.Text.Distilbert, :for_token_classification},
     "GPT2ForSequenceClassification" => {Bumblebee.Text.Gpt2, :for_sequence_classification},
     "GPT2ForTokenClassification" => {Bumblebee.Text.Gpt2, :for_token_classification},
     "GPT2LMHeadModel" => {Bumblebee.Text.Gpt2, :for_causal_language_modeling},
@@ -144,6 +150,7 @@ defmodule Bumblebee do
     "albert" => Bumblebee.Text.AlbertTokenizer,
     "bart" => Bumblebee.Text.BartTokenizer,
     "bert" => Bumblebee.Text.BertTokenizer,
+    "distilbert" => Bumblebee.Text.DistilbertTokenizer,
     "camembert" => Bumblebee.Text.CamembertTokenizer,
     "clip" => Bumblebee.Text.ClipTokenizer,
     "gpt2" => Bumblebee.Text.Gpt2Tokenizer,