test: test we can load real model with safetensors

test/bumblebee/audio/whisper_test.exs

@@ -38,6 +38,42 @@ defmodule Bumblebee.Text.WhisperTest do
       )
     end
 
+    test "base model with safetensors" do
+      assert {:ok, %{model: model, params: params, spec: spec}} =
+               Bumblebee.load_model(
+                 {:hf, "openai/whisper-tiny"},
+                 architecture: :base,
+                 params_filename: "model.safetensors"
+               )
+
+      assert %Bumblebee.Audio.Whisper{architecture: :base} = spec
+
+      input_features = Nx.sin(Nx.iota({1, 3000, 80}, type: :f32))
+      decoder_input_ids = Nx.tensor([[50258, 50259, 50359, 50363]])
+
+      inputs = %{
+        "input_features" => input_features,
+        "decoder_input_ids" => decoder_input_ids
+      }
+
+      outputs = Axon.predict(model, params, inputs)
+
+      assert Nx.shape(outputs.hidden_state) == {1, 4, 384}
+
+      assert_all_close(
+        outputs.hidden_state[[.., .., 1..3]],
+        Nx.tensor([
+          [
+            [9.1349, 0.5695, 8.7758],
+            [0.0160, -7.0785, 1.1313],
+            [6.1074, -2.0481, -1.5687],
+            [5.6247, -10.3924, 7.2008]
+          ]
+        ]),
+        atol: 1.0e-4
+      )
+    end
+
     test "for conditional generation model" do
       assert {:ok, %{model: model, params: params, spec: spec}} =
                Bumblebee.load_model({:hf, "openai/whisper-tiny"})

test/bumblebee_test.exs

@@ -19,5 +19,17 @@ defmodule BumblebeeTest do
 
       assert Enum.sort(Map.keys(params)) == Enum.sort(Map.keys(sharded_params))
     end
+
+    test "supports .safetensors params file" do
+      assert {:ok, %{params: params}} = Bumblebee.load_model({:hf, "openai/whisper-tiny"})
+
+      assert {:ok, %{params: safetensors_params}} =
+               Bumblebee.load_model(
+                 {:hf, "openai/whisper-tiny"},
+                 params_filename: "model.safetensors"
+               )
+
+      assert Enum.sort(Map.keys(params)) == Enum.sort(Map.keys(safetensors_params))
+    end
   end
 end