feat: add mpnet model family

elixir-nx · Nov 5, 2024 · a026833 · a026833
1 parent 9b1a937
commit a026833
Show file tree

Hide file tree

Showing 2 changed files with 460 additions and 0 deletions.
diff --git a/lib/bumblebee.ex b/lib/bumblebee.ex
@@ -170,6 +170,8 @@ defmodule Bumblebee do
     "MistralModel" => {Bumblebee.Text.Mistral, :base},
     "MistralForCausalLM" => {Bumblebee.Text.Mistral, :for_causal_language_modeling},
     "MistralForSequenceClassification" => {Bumblebee.Text.Mistral, :for_sequence_classification},
+    "MPNetModel" => {Bumblebee.Text.MPNet, :base},
+    "MPNetForMaskedLM" => {Bumblebee.Text.MPNet, :for_masked_language_modeling},
     "PhiModel" => {Bumblebee.Text.Phi, :base},
     "PhiForCausalLM" => {Bumblebee.Text.Phi, :for_causal_language_modeling},
     "PhiForSequenceClassification" => {Bumblebee.Text.Phi, :for_sequence_classification},