Add phi model

elixir-nx · Mar 1, 2024 · 9e2441d · 9e2441d
1 parent 4c8bb21
commit 9e2441d
Show file tree

Hide file tree

Showing 5 changed files with 453 additions and 5 deletions.
diff --git a/lib/bumblebee.ex b/lib/bumblebee.ex
@@ -164,6 +164,9 @@ defmodule Bumblebee do
     "MistralModel" => {Bumblebee.Text.Mistral, :base},
     "MistralForCausalLM" => {Bumblebee.Text.Mistral, :for_causal_language_modeling},
     "MistralForSequenceClassification" => {Bumblebee.Text.Mistral, :for_sequence_classification},
+    "PhiModel" => {Bumblebee.Text.Phi, :base},
+    "PhiForCausalLM" => {Bumblebee.Text.Phi, :for_causal_language_modeling},
+    "PhiForSequenceClassification" => {Bumblebee.Text.Phi, :for_sequence_classification},
     "ResNetForImageClassification" => {Bumblebee.Vision.ResNet, :for_image_classification},
     "ResNetModel" => {Bumblebee.Vision.ResNet, :base},
     "RobertaForMaskedLM" => {Bumblebee.Text.Roberta, :for_masked_language_modeling},

diff --git a/lib/bumblebee/layers/transformer.ex b/lib/bumblebee/layers/transformer.ex
@@ -285,11 +285,11 @@ defmodule Bumblebee.Layers.Transformer do
 
         * `:max_positions` - the maximum number of distinct positions
 
-    * `:rotary_embedding_base` - base for computing rotary embedding frequency. Defaults
-      to `10_000`.
+        * `:base` - base for computing rotary embedding frequency. Defaults
+        to `10_000`.
 
-    * `:rotary_percentage` - percentage of hidden dimensions to allocate to rotary embeddings.
-      Defaults to `1.0`.
+        * `:percentage` - percentage of hidden dimensions to allocate to rotary embeddings.
+        Defaults to `1.0`.
 
     * `:name` - the prefix for layer names
 

diff --git a/lib/bumblebee/text/mistral.ex b/lib/bumblebee/text/mistral.ex
@@ -362,7 +362,7 @@ defmodule Bumblebee.Text.Mistral do
 
     gate = Axon.dense(hidden_state, intermediate_size, name: join(name, "gate"), use_bias: false)
 
-    hidden_state = Axon.multiply(intermediate, Axon.activation(gate, activation))
+    hidden_state = Axon.multiply(intermediate, Layers.activation(gate, activation))
 
     Axon.dense(hidden_state, output_size, name: join(name, "output"), use_bias: false)
   end