Add Phi-3 model

lib/bumblebee.ex

@@ -172,6 +172,10 @@ defmodule Bumblebee do
     "PhiForCausalLM" => {Bumblebee.Text.Phi, :for_causal_language_modeling},
     "PhiForSequenceClassification" => {Bumblebee.Text.Phi, :for_sequence_classification},
     "PhiForTokenClassification" => {Bumblebee.Text.Phi, :for_token_classification},
+    "Phi3Model" => {Bumblebee.Text.Phi3, :base},
+    "Phi3ForCausalLM" => {Bumblebee.Text.Phi3, :for_causal_language_modeling},
+    "Phi3ForSequenceClassification" => {Bumblebee.Text.Phi3, :for_sequence_classification},
+    "Phi3ForTokenClassification" => {Bumblebee.Text.Phi3, :for_token_classification},
     "ResNetForImageClassification" => {Bumblebee.Vision.ResNet, :for_image_classification},
     "ResNetModel" => {Bumblebee.Vision.ResNet, :base},
     "RobertaForMaskedLM" => {Bumblebee.Text.Roberta, :for_masked_language_modeling},
@@ -244,6 +248,7 @@ defmodule Bumblebee do
     "mistral" => :llama,
     "mbart" => :mbart,
     "phi" => :code_gen,
+    "phi3" => :llama,
     "roberta" => :roberta,
     "t5" => :t5,
     "whisper" => :whisper,

lib/bumblebee/layers.ex

@@ -1230,25 +1230,66 @@ defmodule Bumblebee.Layers do
                 ) do
     position = Nx.iota({sequence_length})
 
-    {base, position} =
-      case scaling_strategy do
-        %{type: :linear, factor: factor} ->
-          {base, Nx.divide(position, factor)}
-
-        %{type: :dynamic, factor: factor} when sequence_length > max_positions ->
-          base =
-            base
-            |> Nx.multiply(factor * sequence_length / max_positions - (factor - 1))
-            |> Nx.pow(size / (size - 2))
-
-          {base, position}
+    range = Nx.iota({div(size, 2)}) |> Nx.multiply(2) |> Nx.divide(size)
 
-        _other ->
-          {base, position}
-      end
+    case scaling_strategy do
+      %{type: :linear, factor: factor} ->
+        frequency = Nx.pow(base, range)
+        position = Nx.divide(position, factor)
+        positions_cos_sin(position, frequency)
+
+      %{type: :dynamic, factor: factor} when sequence_length > max_positions ->
+        base =
+          base
+          |> Nx.multiply(factor * sequence_length / max_positions - (factor - 1))
+          |> Nx.pow(size / (size - 2))
+
+        frequency = Nx.pow(base, range)
+        positions_cos_sin(position, frequency)
+
+      %{
+        type: type,
+        short_factor: short_factor,
+        long_factor: long_factor,
+        original_max_positions: original_max_positions
+      }
+      when type in [:su, :yarn] ->
+        factor =
+          if sequence_length > original_max_positions do
+            Nx.tensor(long_factor, type: :f32)
+          else
+            Nx.tensor(short_factor, type: :f32)
+          end
+
+        scale = max_positions / original_max_positions
+
+        cos_sin_factor =
+          cond do
+            scale <= 1.0 ->
+              1.0
+
+            type == :su ->
+              Nx.divide(Nx.log(scale), Nx.log(original_max_positions))
+              |> Nx.add(1)
+              |> Nx.sqrt()
+
+            type == :yarn ->
+              Nx.multiply(0.1, Nx.log(scale))
+              |> Nx.add(1.0)
+          end
+
+        frequency = Nx.multiply(factor, Nx.pow(base, range))
+        {cos, sin} = positions_cos_sin(position, frequency)
+        {Nx.multiply(cos, cos_sin_factor), Nx.multiply(sin, cos_sin_factor)}
+
+      _other ->
+        frequency = Nx.pow(base, range)
+        positions_cos_sin(position, frequency)
+    end
+  end
 
-    range = Nx.iota({div(size, 2)}) |> Nx.multiply(2) |> Nx.divide(size)
-    inv_frequency = Nx.divide(1.0, Nx.pow(base, range))
+  defnp positions_cos_sin(position, frequency) do
+    inv_frequency = 1.0 / frequency
     angle = Nx.outer(position, inv_frequency)
 
     angle = Nx.concatenate([angle, angle], axis: -1)