diff --git a/lib/bumblebee.ex b/lib/bumblebee.ex
index 28314c44..adf88a2b 100644
--- a/lib/bumblebee.ex
+++ b/lib/bumblebee.ex
@@ -640,6 +640,80 @@ defmodule Bumblebee do
     end
   end
 
+  @doc """
+  Loads an embedding head from a SentenceTransformers model repository and attaches
+  it to the given model.
+
+  Reads `modules.json` from the repository and stacks the corresponding layers
+  (such as pooling, dense projections, and normalization) on top of `model_info.model`,
+  and loads the corresponding parameters into `model_info.params`.
+
+  The final output of the resulting model is a map with the `:embedding` key.
+
+  Supported SentenceTransformers modules include:
+    * `sentence_transformers.models.Pooling` (modes: `:mean_tokens`, `:cls_token`,
+      `:max_tokens`, `:mean_sqrt_len_tokens`, `:last_token`)
+    * `sentence_transformers.models.Dense` (activations: `:identity`, `:tanh`,
+      `:relu`, `:gelu`, `:silu`, loaded from Safetensors or PyTorch weights)
+    * `sentence_transformers.models.Normalize` (L2 normalization)
+    * `sentence_transformers.models.Dropout` (no-op in inference)
+
+  > #### Note {: .info}
+  >
+  > SentenceTransformers heads attach to the base encoder's hidden states. If a model repository
+  > config defaults to another architecture (such as `:for_masked_language_modeling`), ensure
+  > you specify `architecture: :base` when calling `load_model/2`.
+
+  ## Options
+
+    * `:fuse_dense` - when true, fuses consecutive linear projection layers
+      without bias into a single projection layer. Defaults to `false`
+
+    * `:safetensors_reader` - a function that reads a safetensors file into
+      parameters map. Defaults to `Safetensors.read!/2` with `lazy: true`
+
+    * `:backend` - the backend to allocate the tensors on. It is either
+      an atom or a tuple in the shape `{backend, options}`
+
+    * `:type` - either a type or `Axon.MixedPrecision` policy to apply
+      to the model parameters
+
+  ## Examples
+
+      {:ok, model_info} = Bumblebee.load_model({:hf, "unsloth/embeddinggemma-300m"})
+      {:ok, model_info} = Bumblebee.load_embedding_head({:hf, "unsloth/embeddinggemma-300m"}, model_info, fuse_dense: true)
+      serving = Bumblebee.Text.text_embedding(model_info, tokenizer)
+
+  """
+  @doc type: :model
+  @spec load_embedding_head(repository(), model_info(), keyword()) ::
+          {:ok, model_info()} | {:error, String.t()}
+  def load_embedding_head(repository, model_info, opts \\ []) do
+    repository = normalize_repository!(repository)
+
+    opts =
+      Keyword.validate!(opts, [
+        :safetensors_reader,
+        :backend,
+        :type,
+        fuse_dense: false
+      ])
+
+    case get_repo_files(repository) do
+      {:ok, repo_files} ->
+        HuggingFace.SentenceTransformers.load_embedding_head(
+          repository,
+          repo_files,
+          &download(repository, &1, repo_files[&1]),
+          model_info,
+          opts
+        )
+
+      {:error, message} ->
+        {:error, message}
+    end
+  end
+
   defp maybe_load_model_spec(opts, repository, repo_files) do
     spec_result =
       if spec = opts[:spec] do
@@ -1275,13 +1349,14 @@ defmodule Bumblebee do
 
   defp get_repo_files({:local, dir}) do
     case File.ls(dir) do
-      {:ok, filenames} ->
+      {:ok, _filenames} ->
+        paths = Path.wildcard(Path.join(dir, "**"))
+
         repo_files =
-          for filename <- filenames,
-              path = Path.join(dir, filename),
+          for path <- paths,
               File.regular?(path),
               into: %{},
-              do: {filename, nil}
+              do: {Path.relative_to(path, dir), nil}
 
         {:ok, repo_files}
 
diff --git a/lib/bumblebee/huggingface/hub.ex b/lib/bumblebee/huggingface/hub.ex
index 6a730301..312a2af6 100644
--- a/lib/bumblebee/huggingface/hub.ex
+++ b/lib/bumblebee/huggingface/hub.ex
@@ -21,7 +21,7 @@ defmodule Bumblebee.HuggingFace.Hub do
   def file_listing_url(repository_id, subdir, revision) do
     revision = revision || "main"
     path = if(subdir, do: "/" <> subdir)
-    @huggingface_endpoint <> "/api/models/#{repository_id}/tree/#{revision}#{path}"
+    @huggingface_endpoint <> "/api/models/#{repository_id}/tree/#{revision}#{path}?recursive=true"
   end
 
   @doc """
diff --git a/lib/bumblebee/huggingface/sentence_transformers.ex b/lib/bumblebee/huggingface/sentence_transformers.ex
new file mode 100644
index 00000000..a694a7fb
--- /dev/null
+++ b/lib/bumblebee/huggingface/sentence_transformers.ex
@@ -0,0 +1,506 @@
+defmodule Bumblebee.HuggingFace.SentenceTransformers do
+  @moduledoc false
+
+  alias Bumblebee.Layers
+
+  @doc """
+  Loads a SentenceTransformers embedding head and attaches it to the model.
+  """
+  def load_embedding_head(_repository, repo_files, download_fun, model_info, opts) do
+    opts =
+      opts
+      |> Keyword.put_new_lazy(:type, fn -> infer_param_type(model_info.params) end)
+      |> Keyword.put_new_lazy(:backend, fn -> infer_param_backend(model_info.params) end)
+
+    case repo_files do
+      %{"modules.json" => _etag} ->
+        with {:ok, modules_path} <- download_fun.("modules.json"),
+             {:ok, modules} <- decode_json(modules_path) do
+          modules = Enum.sort_by(modules, &(&1["idx"] || 0))
+
+          case modules do
+            [%{"type" => "sentence_transformers.models.Transformer"} | remaining] ->
+              with {:ok, remaining} <-
+                     maybe_fuse_dense(remaining, opts[:fuse_dense], repo_files, download_fun) do
+                build_pipeline(remaining, repo_files, download_fun, model_info, opts)
+              end
+
+            _other ->
+              {:error,
+               "expected the first module in modules.json to be sentence_transformers.models.Transformer"}
+          end
+        end
+
+      _ ->
+        {:error, "could not find modules.json in the repository"}
+    end
+  end
+
+  defp build_pipeline(modules, repo_files, download_fun, model_info, opts) do
+    base_model = model_info.model
+
+    hidden_state =
+      Axon.nx(base_model, fn
+        %{hidden_state: hidden_state} ->
+          hidden_state
+
+        output ->
+          raise ArgumentError,
+                "expected model output to contain :hidden_state, but got keys: #{inspect(Map.keys(output))}." <>
+                  " If the model defaults to a language modeling architecture, please specify" <>
+                  " `architecture: :base` when calling `Bumblebee.load_model/2`."
+      end)
+
+    attention_mask =
+      Layers.default Axon.input("attention_mask", optional: true) do
+        Layers.default_attention_mask(Axon.input("input_ids"))
+      end
+
+    initial_state = {hidden_state, model_info.params.data, model_info.params.parameters}
+
+    result =
+      Enum.reduce_while(
+        modules,
+        {:ok, initial_state},
+        fn module, {:ok, {node, params_data, params_parameters}} ->
+          case build_module(
+                 module,
+                 node,
+                 attention_mask,
+                 params_data,
+                 params_parameters,
+                 repo_files,
+                 download_fun,
+                 opts
+               ) do
+            {:ok, node, params_data, params_parameters} ->
+              {:cont, {:ok, {node, params_data, params_parameters}}}
+
+            {:error, reason} ->
+              {:halt, {:error, reason}}
+          end
+        end
+      )
+
+    with {:ok, {embedding, params_data, params_parameters}} <- result do
+      final_model =
+        Axon.container(%{
+          embedding: embedding,
+          pooled_state: embedding,
+          hidden_state: hidden_state
+        })
+
+      final_model = apply_type(final_model, opts[:type])
+
+      final_params =
+        if params_parameters do
+          %{model_info.params | data: params_data, parameters: params_parameters}
+        else
+          %{model_info.params | data: params_data}
+        end
+
+      {:ok, %{model_info | model: final_model, params: final_params}}
+    end
+  end
+
+  defp build_module(
+         %{"type" => "sentence_transformers.models.Pooling", "path" => path},
+         node,
+         attention_mask,
+         params_data,
+         params_parameters,
+         _repo_files,
+         download_fun,
+         _opts
+       ) do
+    config_file = Path.join(path, "config.json")
+
+    with {:ok, config_path} <- download_fun.(config_file),
+         {:ok, config} <- decode_json(config_path),
+         {:ok, node} <- pooling_layer(node, attention_mask, config, path) do
+      {:ok, node, params_data, params_parameters}
+    end
+  end
+
+  defp build_module(
+         %{"type" => "sentence_transformers.models.Dense", "path" => path},
+         node,
+         _attention_mask,
+         params_data,
+         params_parameters,
+         repo_files,
+         download_fun,
+         opts
+       ) do
+    config_file = Path.join(path, "config.json")
+
+    with {:ok, config_path} <- download_fun.(config_file),
+         {:ok, config} <- decode_json(config_path),
+         {:ok, activation} <- parse_activation(config["activation_function"], path),
+         {:ok, weights_file} <- find_weights_file(repo_files, path),
+         {:ok, weights_path} <- download_fun.(weights_file) do
+      out_features = config["out_features"]
+      bias = Map.get(config, "bias", true)
+
+      layer_name = path
+
+      node =
+        node
+        |> Axon.dense(out_features, use_bias: bias, name: layer_name)
+        |> maybe_activation(activation)
+
+      tensors = load_tensors(weights_file, weights_path, opts)
+
+      kernel =
+        tensors["linear.weight"]
+        |> Nx.to_tensor()
+        |> Nx.transpose()
+        |> cast_param(opts[:type])
+        |> allocate_param(opts[:backend])
+
+      layer_params =
+        if bias do
+          bias_tensor =
+            tensors["linear.bias"]
+            |> Nx.to_tensor()
+            |> cast_param(opts[:type])
+            |> allocate_param(opts[:backend])
+
+          %{"kernel" => kernel, "bias" => bias_tensor}
+        else
+          %{"kernel" => kernel}
+        end
+
+      params_data = Map.put(params_data, layer_name, layer_params)
+
+      params_parameters =
+        if params_parameters do
+          Map.put(params_parameters, layer_name, Map.keys(layer_params))
+        end
+
+      {:ok, node, params_data, params_parameters}
+    end
+  end
+
+  defp build_module(
+         %{
+           "type" => "sentence_transformers.models.FusedDense",
+           "path" => layer_name,
+           "paths" => [p1, p2],
+           "out_features" => out_features
+         },
+         node,
+         _attention_mask,
+         params_data,
+         params_parameters,
+         repo_files,
+         download_fun,
+         opts
+       ) do
+    with {:ok, w1_file} <- find_weights_file(repo_files, p1),
+         {:ok, w1_path} <- download_fun.(w1_file),
+         {:ok, w2_file} <- find_weights_file(repo_files, p2),
+         {:ok, w2_path} <- download_fun.(w2_file) do
+      t1 = load_tensors(w1_file, w1_path, opts)
+      t2 = load_tensors(w2_file, w2_path, opts)
+
+      k1 = t1["linear.weight"] |> Nx.to_tensor() |> Nx.transpose()
+      k2 = t2["linear.weight"] |> Nx.to_tensor() |> Nx.transpose()
+
+      k_fused =
+        Nx.dot(k1, k2)
+        |> cast_param(opts[:type])
+        |> allocate_param(opts[:backend])
+
+      node = Axon.dense(node, out_features, use_bias: false, name: layer_name)
+      params_data = Map.put(params_data, layer_name, %{"kernel" => k_fused})
+
+      params_parameters =
+        if params_parameters do
+          Map.put(params_parameters, layer_name, ["kernel"])
+        end
+
+      {:ok, node, params_data, params_parameters}
+    end
+  end
+
+  defp build_module(
+         %{"type" => "sentence_transformers.models.Normalize", "path" => path},
+         node,
+         _attention_mask,
+         params_data,
+         params_parameters,
+         _repo_files,
+         _download_fun,
+         _opts
+       ) do
+    node = Axon.nx(node, &Bumblebee.Utils.Nx.normalize/1, name: "#{path}.normalize")
+    {:ok, node, params_data, params_parameters}
+  end
+
+  defp build_module(
+         %{"type" => "sentence_transformers.models.Dropout"},
+         node,
+         _attention_mask,
+         params_data,
+         params_parameters,
+         _repo_files,
+         _download_fun,
+         _opts
+       ) do
+    {:ok, node, params_data, params_parameters}
+  end
+
+  defp build_module(
+         %{"type" => type},
+         _node,
+         _mask,
+         _params,
+         _params_parameters,
+         _repo_files,
+         _download_fun,
+         _opts
+       ) do
+    {:error, "unsupported SentenceTransformers module #{inspect(type)}"}
+  end
+
+  defp pooling_layer(hidden_state, attention_mask, config, name) do
+    modes =
+      for {keys, mode} <- [
+            {["pooling_mode_cls_token"], :cls_token},
+            {["pooling_mode_max_tokens"], :max_tokens},
+            {["pooling_mode_mean_tokens"], :mean_tokens},
+            {["pooling_mode_mean_sqrt_len_tokens"], :mean_sqrt_len_tokens},
+            {["pooling_mode_lasttoken", "pooling_mode_last_token"], :last_token}
+          ],
+          Enum.any?(keys, &(config[&1] == true)),
+          do: mode
+
+    if modes == [] do
+      {:error, "no supported pooling mode found in #{name}"}
+    else
+      layer =
+        Axon.layer(
+          fn hidden_state, attention_mask, _opts ->
+            type = Nx.type(hidden_state)
+
+            eps =
+              case type do
+                {:f, 16} -> Nx.tensor(1.0e-4, type: type)
+                {:bf, 16} -> Nx.tensor(1.0e-4, type: type)
+                _ -> Nx.tensor(1.0e-9, type: type)
+              end
+
+            pool_outputs =
+              Enum.map(modes, fn
+                :mean_tokens ->
+                  mask = attention_mask |> Nx.as_type(type) |> Nx.new_axis(-1)
+                  sum_embeddings = Nx.sum(Nx.multiply(hidden_state, mask), axes: [1])
+                  sum_mask = Nx.sum(mask, axes: [1]) |> Nx.max(eps)
+                  Nx.divide(sum_embeddings, sum_mask)
+
+                :cls_token ->
+                  hidden_state[[.., 0, ..]]
+
+                :max_tokens ->
+                  mask = Nx.new_axis(attention_mask, -1)
+                  pred = Nx.broadcast(Nx.not_equal(mask, 0), Nx.shape(hidden_state))
+
+                  pred
+                  |> Nx.select(hidden_state, Nx.Constants.min_finite(type))
+                  |> Nx.reduce_max(axes: [1])
+
+                :mean_sqrt_len_tokens ->
+                  mask = attention_mask |> Nx.as_type(type) |> Nx.new_axis(-1)
+                  sum_embeddings = Nx.sum(Nx.multiply(hidden_state, mask), axes: [1])
+                  sum_mask = Nx.sum(mask, axes: [1]) |> Nx.max(eps)
+                  Nx.divide(sum_embeddings, Nx.sqrt(sum_mask))
+
+                :last_token ->
+                  lengths =
+                    attention_mask
+                    |> Nx.sum(axes: [1])
+                    |> Nx.subtract(1)
+                    |> Nx.as_type({:s, 64})
+
+                  Bumblebee.Utils.Nx.batched_take(hidden_state, lengths)
+              end)
+
+            case pool_outputs do
+              [single] -> single
+              multiple -> Nx.concatenate(multiple, axis: -1)
+            end
+          end,
+          [hidden_state, attention_mask],
+          name: name
+        )
+
+      {:ok, layer}
+    end
+  end
+
+  defp parse_activation("torch.nn.modules.linear.Identity", _path), do: {:ok, nil}
+  defp parse_activation("torch.nn.modules.activation.Tanh", _path), do: {:ok, :tanh}
+  defp parse_activation("torch.nn.modules.activation.ReLU", _path), do: {:ok, :relu}
+  defp parse_activation("torch.nn.modules.activation.GELU", _path), do: {:ok, :gelu}
+  defp parse_activation("torch.nn.modules.activation.SiLU", _path), do: {:ok, :silu}
+  defp parse_activation(nil, _path), do: {:ok, nil}
+
+  defp parse_activation(other, path) do
+    {:error, "unsupported activation function #{inspect(other)} in #{path}"}
+  end
+
+  defp maybe_activation(node, nil), do: node
+  defp maybe_activation(node, activation), do: Axon.activation(node, activation)
+
+  defp maybe_fuse_dense(modules, true, repo_files, download_fun) do
+    fuse_dense_modules(modules, repo_files, download_fun, [])
+  end
+
+  defp maybe_fuse_dense(modules, _false, _repo_files, _download_fun), do: {:ok, modules}
+
+  defp fuse_dense_modules(
+         [
+           %{"type" => "sentence_transformers.models.Dense", "path" => p1} = m1,
+           %{"type" => "sentence_transformers.models.Dense", "path" => p2} = m2 | rest
+         ],
+         repo_files,
+         download_fun,
+         acc
+       ) do
+    with {:ok, c1} <- load_module_config(p1, download_fun),
+         {:ok, c2} <- load_module_config(p2, download_fun),
+         true <- can_fuse_dense?(c1, c2) do
+      fused_module = %{
+        "type" => "sentence_transformers.models.FusedDense",
+        "path" => "#{p1}_#{p2}",
+        "paths" => [p1, p2],
+        "out_features" => c2["out_features"]
+      }
+
+      fuse_dense_modules(rest, repo_files, download_fun, [fused_module | acc])
+    else
+      _ ->
+        fuse_dense_modules([m2 | rest], repo_files, download_fun, [m1 | acc])
+    end
+  end
+
+  defp fuse_dense_modules([other | rest], repo_files, download_fun, acc) do
+    fuse_dense_modules(rest, repo_files, download_fun, [other | acc])
+  end
+
+  defp fuse_dense_modules([], _repo_files, _download_fun, acc) do
+    {:ok, Enum.reverse(acc)}
+  end
+
+  defp load_module_config(path, download_fun) do
+    config_file = Path.join(path, "config.json")
+
+    with {:ok, config_path} <- download_fun.(config_file),
+         {:ok, config} <- decode_json(config_path) do
+      {:ok, config}
+    end
+  end
+
+  defp can_fuse_dense?(c1, c2) do
+    c1_linear? = c1["activation_function"] in [nil, "torch.nn.modules.linear.Identity"]
+    c2_linear? = c2["activation_function"] in [nil, "torch.nn.modules.linear.Identity"]
+    no_bias? = !Map.get(c1, "bias", true) and !Map.get(c2, "bias", true)
+    dims_match? = c1["out_features"] == c2["in_features"]
+
+    c1_linear? and c2_linear? and no_bias? and dims_match?
+  end
+
+  defp find_weights_file(repo_files, dir) do
+    safetensors = Path.join(dir, "model.safetensors")
+    pytorch_bin = Path.join(dir, "pytorch_model.bin")
+
+    cond do
+      Map.has_key?(repo_files, safetensors) ->
+        {:ok, safetensors}
+
+      Map.has_key?(repo_files, pytorch_bin) ->
+        {:ok, pytorch_bin}
+
+      true ->
+        {:error, "could not find parameters file in #{dir}"}
+    end
+  end
+
+  defp load_tensors(weights_file, weights_path, opts) do
+    case Path.extname(weights_file) do
+      ".safetensors" ->
+        reader = opts[:safetensors_reader] || (&Safetensors.read!(&1, lazy: true))
+        reader.(weights_path)
+
+      _ ->
+        Bumblebee.Conversion.PyTorchLoader.load!(weights_path)
+    end
+  end
+
+  defp cast_param(tensor, nil), do: tensor
+
+  defp cast_param(tensor, %Axon.MixedPrecision.Policy{params: type}) do
+    Nx.as_type(tensor, type)
+  end
+
+  defp cast_param(tensor, type) do
+    type = Nx.Type.normalize!(type)
+    Nx.as_type(tensor, type)
+  end
+
+  defp allocate_param(tensor, nil), do: tensor
+
+  defp allocate_param(tensor, backend) do
+    Nx.with_default_backend(backend, fn -> Nx.backend_copy(tensor) end)
+  end
+
+  defp apply_type(model, nil), do: model
+
+  defp apply_type(model, %Axon.MixedPrecision.Policy{} = policy) do
+    Axon.MixedPrecision.apply_policy(model, policy)
+  end
+
+  defp apply_type(model, type) do
+    type = Nx.Type.normalize!(type)
+    policy = Axon.MixedPrecision.create_policy(params: type, compute: type, output: type)
+    Axon.MixedPrecision.apply_policy(model, policy)
+  end
+
+  defp infer_param_type(%Axon.ModelState{data: data}) do
+    find_first_tensor(data, &Nx.type/1)
+  end
+
+  defp infer_param_type(_), do: nil
+
+  defp infer_param_backend(%Axon.ModelState{data: data}) do
+    find_first_tensor(data, fn %Nx.Tensor{data: %backend{}} ->
+      if backend == Nx.BinaryBackend, do: nil, else: backend
+    end)
+  end
+
+  defp infer_param_backend(_), do: nil
+
+  defp find_first_tensor(data, fun) when is_map(data) do
+    Enum.find_value(data, fn
+      {_key, %Nx.Tensor{} = tensor} -> fun.(tensor)
+      {_key, nested} when is_map(nested) -> find_first_tensor(nested, fun)
+      _ -> nil
+    end)
+  end
+
+  defp find_first_tensor(_, _fun), do: nil
+
+  defp decode_json(path) do
+    case File.read(path) do
+      {:ok, content} ->
+        case Jason.decode(content) do
+          {:ok, data} -> {:ok, data}
+          _ -> {:error, "failed to parse #{path} as JSON"}
+        end
+
+      {:error, reason} ->
+        {:error, "failed to read #{path}, reason: #{:file.format_error(reason)}"}
+    end
+  end
+end
diff --git a/lib/bumblebee/text/text_embedding.ex b/lib/bumblebee/text/text_embedding.ex
index 0e2f3278..9d123b6c 100644
--- a/lib/bumblebee/text/text_embedding.ex
+++ b/lib/bumblebee/text/text_embedding.ex
@@ -53,6 +53,9 @@ defmodule Bumblebee.Text.TextEmbedding do
           %{^output_attribute => output} ->
             output
 
+          %{embedding: output} when output_attribute == :pooled_state ->
+            output
+
           %{} ->
             keys = output |> Map.keys() |> Enum.sort()
 
diff --git a/test/bumblebee/huggingface/sentence_transformers_test.exs b/test/bumblebee/huggingface/sentence_transformers_test.exs
new file mode 100644
index 00000000..44435b08
--- /dev/null
+++ b/test/bumblebee/huggingface/sentence_transformers_test.exs
@@ -0,0 +1,1161 @@
+defmodule Bumblebee.HuggingFace.SentenceTransformersTest do
+  use ExUnit.Case, async: true
+
+  import Bumblebee.TestHelpers
+
+  @moduletag model_test_tags()
+
+  setup do
+    model =
+      Axon.input("input_ids", shape: {nil, nil})
+      |> Axon.nx(fn input_ids ->
+        # Mock base model output with hidden_state {batch_size, seq_len, 4}
+        batch_size = Nx.axis_size(input_ids, 0)
+        seq_len = Nx.axis_size(input_ids, 1)
+
+        hidden_state =
+          Nx.broadcast(1.0, {batch_size, seq_len, 4})
+          |> Nx.as_type({:f, 32})
+
+        %{hidden_state: hidden_state}
+      end)
+
+    params = Axon.ModelState.empty()
+    spec = Bumblebee.configure(Bumblebee.Text.Bert, architecture: :base)
+
+    model_info = %{model: model, params: params, spec: spec}
+
+    [model_info: model_info, spec: spec]
+  end
+
+  describe "load_embedding_head/3" do
+    @tag :tmp_dir
+    test "loads pooling and normalize head", %{model_info: model_info, tmp_dir: dir} do
+      modules = [
+        %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"},
+        %{"idx" => 1, "path" => "1_Pooling", "type" => "sentence_transformers.models.Pooling"},
+        %{"idx" => 2, "path" => "2_Normalize", "type" => "sentence_transformers.models.Normalize"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      pooling_dir = Path.join(dir, "1_Pooling")
+      File.mkdir_p!(pooling_dir)
+
+      pooling_config = %{
+        "word_embedding_dimension" => 4,
+        "pooling_mode_cls_token" => false,
+        "pooling_mode_mean_tokens" => true,
+        "pooling_mode_max_tokens" => false
+      }
+
+      File.write!(Path.join(pooling_dir, "config.json"), Jason.encode!(pooling_config))
+
+      assert {:ok, new_model_info} = Bumblebee.load_embedding_head({:local, dir}, model_info)
+
+      inputs = %{
+        "input_ids" => Nx.tensor([[1, 2, 3]]),
+        "attention_mask" => Nx.tensor([[1, 1, 0]])
+      }
+
+      {_init_fn, predict_fn} = Axon.build(new_model_info.model)
+      output = predict_fn.(new_model_info.params, inputs)
+
+      assert Map.has_key?(output, :embedding)
+      assert Map.has_key?(output, :pooled_state)
+      assert Map.has_key?(output, :hidden_state)
+      assert Nx.shape(output.embedding) == {1, 4}
+      assert_all_close(output.pooled_state, output.embedding)
+
+      # Normalization check: L2 norm of output embedding must be 1.0
+      norm = Nx.LinAlg.norm(output.embedding, axes: [-1])
+      assert_all_close(norm, Nx.tensor([1.0]), atol: 1.0e-5)
+    end
+
+    @tag :tmp_dir
+    test "loads dense projection layers with parameters", %{model_info: model_info, tmp_dir: dir} do
+      modules = [
+        %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"},
+        %{"idx" => 1, "path" => "1_Pooling", "type" => "sentence_transformers.models.Pooling"},
+        %{"idx" => 2, "path" => "2_Dense", "type" => "sentence_transformers.models.Dense"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      pooling_dir = Path.join(dir, "1_Pooling")
+      File.mkdir_p!(pooling_dir)
+
+      File.write!(
+        Path.join(pooling_dir, "config.json"),
+        Jason.encode!(%{"pooling_mode_mean_tokens" => true})
+      )
+
+      dense_dir = Path.join(dir, "2_Dense")
+      File.mkdir_p!(dense_dir)
+
+      dense_config = %{
+        "in_features" => 4,
+        "out_features" => 8,
+        "bias" => true,
+        "activation_function" => "torch.nn.modules.linear.Identity"
+      }
+
+      File.write!(Path.join(dense_dir, "config.json"), Jason.encode!(dense_config))
+
+      # Weight in PyTorch format: {out_features, in_features} = {8, 4}
+      weights = %{
+        "linear.weight" => Nx.broadcast(0.5, {8, 4}) |> Nx.as_type({:f, 32}),
+        "linear.bias" => Nx.broadcast(0.1, {8}) |> Nx.as_type({:f, 32})
+      }
+
+      Safetensors.write!(Path.join(dense_dir, "model.safetensors"), weights)
+
+      assert {:ok, new_model_info} = Bumblebee.load_embedding_head({:local, dir}, model_info)
+
+      assert Map.has_key?(new_model_info.params.data, "2_Dense")
+      assert Nx.shape(new_model_info.params.data["2_Dense"]["kernel"]) == {4, 8}
+      assert Nx.shape(new_model_info.params.data["2_Dense"]["bias"]) == {8}
+
+      inputs = %{
+        "input_ids" => Nx.tensor([[1, 2, 3]]),
+        "attention_mask" => Nx.tensor([[1, 1, 1]])
+      }
+
+      {_init_fn, predict_fn} = Axon.build(new_model_info.model)
+      output = predict_fn.(new_model_info.params, inputs)
+
+      assert Nx.shape(output.embedding) == {1, 8}
+    end
+
+    @tag :tmp_dir
+    test "updates ModelState.parameters for added layers", %{spec: spec, tmp_dir: dir} do
+      base_model =
+        Axon.input("input_ids", shape: {nil, 4})
+        |> Axon.nx(fn x -> %{hidden_state: Nx.new_axis(x, 1)} end)
+
+      base_params = %Axon.ModelState{
+        data: %{"base" => %{"kernel" => Nx.broadcast(1.0, {4, 4})}},
+        parameters: %{"base" => ["kernel"]},
+        frozen_parameters: %{},
+        state: %{}
+      }
+
+      model_info = %{model: base_model, params: base_params, spec: spec}
+
+      modules = [
+        %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"},
+        %{"idx" => 1, "path" => "1_Pooling", "type" => "sentence_transformers.models.Pooling"},
+        %{"idx" => 2, "path" => "2_Dense", "type" => "sentence_transformers.models.Dense"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      pooling_dir = Path.join(dir, "1_Pooling")
+      File.mkdir_p!(pooling_dir)
+
+      File.write!(
+        Path.join(pooling_dir, "config.json"),
+        Jason.encode!(%{"pooling_mode_mean_tokens" => true})
+      )
+
+      dense_dir = Path.join(dir, "2_Dense")
+      File.mkdir_p!(dense_dir)
+
+      dense_config = %{
+        "in_features" => 4,
+        "out_features" => 8,
+        "bias" => true,
+        "activation_function" => "torch.nn.modules.linear.Identity"
+      }
+
+      File.write!(Path.join(dense_dir, "config.json"), Jason.encode!(dense_config))
+
+      weights = %{
+        "linear.weight" => Nx.broadcast(0.5, {8, 4}) |> Nx.as_type({:f, 32}),
+        "linear.bias" => Nx.broadcast(0.1, {8}) |> Nx.as_type({:f, 32})
+      }
+
+      Safetensors.write!(Path.join(dense_dir, "model.safetensors"), weights)
+
+      assert {:ok, new_model_info} = Bumblebee.load_embedding_head({:local, dir}, model_info)
+
+      assert Enum.sort(new_model_info.params.parameters["2_Dense"]) == ["bias", "kernel"]
+      assert new_model_info.params.parameters["base"] == ["kernel"]
+    end
+
+    @tag :tmp_dir
+    test "infers type from base model params when type option is omitted", %{
+      spec: spec,
+      tmp_dir: dir
+    } do
+      base_model =
+        Axon.input("input_ids", shape: {nil, nil})
+        |> Axon.nx(fn input_ids ->
+          batch_size = Nx.axis_size(input_ids, 0)
+          seq_len = Nx.axis_size(input_ids, 1)
+
+          hidden_state =
+            Nx.broadcast(1.0, {batch_size, seq_len, 4})
+            |> Nx.as_type({:f, 16})
+
+          %{hidden_state: hidden_state}
+        end)
+
+      base_params = %Axon.ModelState{
+        data: %{"base" => %{"kernel" => Nx.broadcast(1.0, {4, 4}) |> Nx.as_type({:f, 16})}}
+      }
+
+      model_info = %{model: base_model, params: base_params, spec: spec}
+
+      modules = [
+        %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"},
+        %{"idx" => 1, "path" => "1_Pooling", "type" => "sentence_transformers.models.Pooling"},
+        %{"idx" => 2, "path" => "2_Dense", "type" => "sentence_transformers.models.Dense"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      pooling_dir = Path.join(dir, "1_Pooling")
+      File.mkdir_p!(pooling_dir)
+
+      File.write!(
+        Path.join(pooling_dir, "config.json"),
+        Jason.encode!(%{"pooling_mode_mean_tokens" => true})
+      )
+
+      dense_dir = Path.join(dir, "2_Dense")
+      File.mkdir_p!(dense_dir)
+
+      dense_config = %{
+        "in_features" => 4,
+        "out_features" => 8,
+        "bias" => true,
+        "activation_function" => "torch.nn.modules.linear.Identity"
+      }
+
+      File.write!(Path.join(dense_dir, "config.json"), Jason.encode!(dense_config))
+
+      weights = %{
+        "linear.weight" => Nx.broadcast(0.5, {8, 4}) |> Nx.as_type({:f, 32}),
+        "linear.bias" => Nx.broadcast(0.1, {8}) |> Nx.as_type({:f, 32})
+      }
+
+      Safetensors.write!(Path.join(dense_dir, "model.safetensors"), weights)
+
+      assert {:ok, new_model_info} = Bumblebee.load_embedding_head({:local, dir}, model_info)
+
+      assert Nx.type(new_model_info.params.data["2_Dense"]["kernel"]) == {:f, 16}
+      assert Nx.type(new_model_info.params.data["2_Dense"]["bias"]) == {:f, 16}
+    end
+
+    @tag :tmp_dir
+    test "returns error when modules.json is missing", %{model_info: model_info, tmp_dir: dir} do
+      assert {:error, "could not find modules.json in the repository"} =
+               Bumblebee.load_embedding_head({:local, dir}, model_info)
+    end
+
+    @tag :tmp_dir
+    test "returns error when first module is not Transformer", %{
+      model_info: model_info,
+      tmp_dir: dir
+    } do
+      modules = [
+        %{"idx" => 0, "path" => "1_Pooling", "type" => "sentence_transformers.models.Pooling"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      assert {:error,
+              "expected the first module in modules.json to be sentence_transformers.models.Transformer"} =
+               Bumblebee.load_embedding_head({:local, dir}, model_info)
+    end
+
+    @tag :tmp_dir
+    test "raises ArgumentError when base model output lacks hidden_state", %{
+      spec: spec,
+      tmp_dir: dir
+    } do
+      lm_model =
+        Axon.input("input_ids", shape: {nil, nil})
+        |> Axon.nx(fn input_ids ->
+          %{
+            logits:
+              Nx.broadcast(0.0, {Nx.axis_size(input_ids, 0), Nx.axis_size(input_ids, 1), 100})
+          }
+        end)
+
+      model_info = %{model: lm_model, params: Axon.ModelState.empty(), spec: spec}
+
+      modules = [
+        %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"},
+        %{"idx" => 1, "path" => "1_Pooling", "type" => "sentence_transformers.models.Pooling"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      pooling_dir = Path.join(dir, "1_Pooling")
+      File.mkdir_p!(pooling_dir)
+
+      File.write!(
+        Path.join(pooling_dir, "config.json"),
+        Jason.encode!(%{"pooling_mode_mean_tokens" => true})
+      )
+
+      assert {:ok, new_model_info} = Bumblebee.load_embedding_head({:local, dir}, model_info)
+
+      {_init_fn, predict_fn} = Axon.build(new_model_info.model)
+
+      assert_raise Axon.CompileError,
+                   ~r/expected model output to contain :hidden_state, but got keys: \[:logits\]/,
+                   fn ->
+                     predict_fn.(new_model_info.params, %{"input_ids" => Nx.tensor([[1, 2, 3]])})
+                   end
+    end
+
+    @tag :tmp_dir
+    test "fuses consecutive linear dense layers when fuse_dense: true", %{
+      model_info: model_info,
+      tmp_dir: dir
+    } do
+      modules = [
+        %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"},
+        %{"idx" => 1, "path" => "1_Pooling", "type" => "sentence_transformers.models.Pooling"},
+        %{"idx" => 2, "path" => "2_Dense", "type" => "sentence_transformers.models.Dense"},
+        %{"idx" => 3, "path" => "3_Dense", "type" => "sentence_transformers.models.Dense"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      pooling_dir = Path.join(dir, "1_Pooling")
+      File.mkdir_p!(pooling_dir)
+
+      File.write!(
+        Path.join(pooling_dir, "config.json"),
+        Jason.encode!(%{"pooling_mode_mean_tokens" => true})
+      )
+
+      # 2_Dense: 4 -> 8
+      dense1_dir = Path.join(dir, "2_Dense")
+      File.mkdir_p!(dense1_dir)
+
+      File.write!(
+        Path.join(dense1_dir, "config.json"),
+        Jason.encode!(%{
+          "in_features" => 4,
+          "out_features" => 8,
+          "bias" => false,
+          "activation_function" => "torch.nn.modules.linear.Identity"
+        })
+      )
+
+      k1 = Nx.broadcast(0.5, {8, 4}) |> Nx.as_type({:f, 32})
+      Safetensors.write!(Path.join(dense1_dir, "model.safetensors"), %{"linear.weight" => k1})
+
+      # 3_Dense: 8 -> 4
+      dense2_dir = Path.join(dir, "3_Dense")
+      File.mkdir_p!(dense2_dir)
+
+      File.write!(
+        Path.join(dense2_dir, "config.json"),
+        Jason.encode!(%{
+          "in_features" => 8,
+          "out_features" => 4,
+          "bias" => false,
+          "activation_function" => "torch.nn.modules.linear.Identity"
+        })
+      )
+
+      k2 = Nx.broadcast(0.25, {4, 8}) |> Nx.as_type({:f, 32})
+      Safetensors.write!(Path.join(dense2_dir, "model.safetensors"), %{"linear.weight" => k2})
+
+      inputs = %{
+        "input_ids" => Nx.tensor([[1, 2, 3]]),
+        "attention_mask" => Nx.tensor([[1, 1, 1]])
+      }
+
+      # Separate (fuse_dense: false)
+      assert {:ok, sep_model_info} =
+               Bumblebee.load_embedding_head({:local, dir}, model_info, fuse_dense: false)
+
+      assert Map.has_key?(sep_model_info.params.data, "2_Dense")
+      assert Map.has_key?(sep_model_info.params.data, "3_Dense")
+
+      {_, predict_sep} = Axon.build(sep_model_info.model)
+      out_sep = predict_sep.(sep_model_info.params, inputs)
+
+      # Fused (fuse_dense: true)
+      assert {:ok, fus_model_info} =
+               Bumblebee.load_embedding_head({:local, dir}, model_info, fuse_dense: true)
+
+      assert Map.has_key?(fus_model_info.params.data, "2_Dense_3_Dense")
+      assert Nx.shape(fus_model_info.params.data["2_Dense_3_Dense"]["kernel"]) == {4, 4}
+
+      {_, predict_fus} = Axon.build(fus_model_info.model)
+      out_fus = predict_fus.(fus_model_info.params, inputs)
+
+      assert_all_close(out_sep.embedding, out_fus.embedding)
+    end
+
+    @tag :tmp_dir
+    test "does not fuse linear dense layers when bias is omitted (defaults to true)", %{
+      model_info: model_info,
+      tmp_dir: dir
+    } do
+      modules = [
+        %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"},
+        %{"idx" => 1, "path" => "1_Pooling", "type" => "sentence_transformers.models.Pooling"},
+        %{"idx" => 2, "path" => "2_Dense", "type" => "sentence_transformers.models.Dense"},
+        %{"idx" => 3, "path" => "3_Dense", "type" => "sentence_transformers.models.Dense"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      pooling_dir = Path.join(dir, "1_Pooling")
+      File.mkdir_p!(pooling_dir)
+
+      File.write!(
+        Path.join(pooling_dir, "config.json"),
+        Jason.encode!(%{"pooling_mode_mean_tokens" => true})
+      )
+
+      # 2_Dense: 4 -> 8 (bias omitted, defaults to true)
+      dense1_dir = Path.join(dir, "2_Dense")
+      File.mkdir_p!(dense1_dir)
+
+      File.write!(
+        Path.join(dense1_dir, "config.json"),
+        Jason.encode!(%{
+          "in_features" => 4,
+          "out_features" => 8,
+          "activation_function" => "torch.nn.modules.linear.Identity"
+        })
+      )
+
+      k1 = Nx.broadcast(0.5, {8, 4}) |> Nx.as_type({:f, 32})
+      b1 = Nx.broadcast(0.1, {8}) |> Nx.as_type({:f, 32})
+
+      Safetensors.write!(Path.join(dense1_dir, "model.safetensors"), %{
+        "linear.weight" => k1,
+        "linear.bias" => b1
+      })
+
+      # 3_Dense: 8 -> 4 (bias omitted, defaults to true)
+      dense2_dir = Path.join(dir, "3_Dense")
+      File.mkdir_p!(dense2_dir)
+
+      File.write!(
+        Path.join(dense2_dir, "config.json"),
+        Jason.encode!(%{
+          "in_features" => 8,
+          "out_features" => 4,
+          "activation_function" => "torch.nn.modules.linear.Identity"
+        })
+      )
+
+      k2 = Nx.broadcast(0.25, {4, 8}) |> Nx.as_type({:f, 32})
+      b2 = Nx.broadcast(0.2, {4}) |> Nx.as_type({:f, 32})
+
+      Safetensors.write!(Path.join(dense2_dir, "model.safetensors"), %{
+        "linear.weight" => k2,
+        "linear.bias" => b2
+      })
+
+      assert {:ok, new_model_info} =
+               Bumblebee.load_embedding_head({:local, dir}, model_info, fuse_dense: true)
+
+      assert Map.has_key?(new_model_info.params.data, "2_Dense")
+      assert Map.has_key?(new_model_info.params.data, "3_Dense")
+      refute Map.has_key?(new_model_info.params.data, "2_Dense_3_Dense")
+    end
+
+    @tag :tmp_dir
+    test "returns error on unsupported module", %{model_info: model_info, tmp_dir: dir} do
+      modules = [
+        %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"},
+        %{"idx" => 1, "path" => "1_Unknown", "type" => "sentence_transformers.models.Unknown"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      assert {:error,
+              "unsupported SentenceTransformers module \"sentence_transformers.models.Unknown\""} =
+               Bumblebee.load_embedding_head({:local, dir}, model_info)
+    end
+
+    @tag :tmp_dir
+    test "loads max_tokens pooling", %{spec: spec, tmp_dir: dir} do
+      custom_model =
+        Axon.input("input_ids", shape: {nil, nil})
+        |> Axon.nx(fn _input_ids ->
+          hidden_state =
+            Nx.tensor([[[1.0, 5.0, 2.0, 3.0], [4.0, 2.0, 6.0, 1.0], [99.0, 99.0, 99.0, 99.0]]])
+
+          %{hidden_state: hidden_state}
+        end)
+
+      model_info = %{model: custom_model, params: Axon.ModelState.empty(), spec: spec}
+
+      modules = [
+        %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"},
+        %{"idx" => 1, "path" => "1_Pooling", "type" => "sentence_transformers.models.Pooling"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      pooling_dir = Path.join(dir, "1_Pooling")
+      File.mkdir_p!(pooling_dir)
+
+      pooling_config = %{
+        "pooling_mode_cls_token" => false,
+        "pooling_mode_mean_tokens" => false,
+        "pooling_mode_max_tokens" => true
+      }
+
+      File.write!(Path.join(pooling_dir, "config.json"), Jason.encode!(pooling_config))
+
+      assert {:ok, new_model_info} = Bumblebee.load_embedding_head({:local, dir}, model_info)
+
+      inputs = %{
+        "input_ids" => Nx.tensor([[1, 2, 3]]),
+        "attention_mask" => Nx.tensor([[1, 1, 0]])
+      }
+
+      {_init_fn, predict_fn} = Axon.build(new_model_info.model)
+      output = predict_fn.(new_model_info.params, inputs)
+
+      assert_all_close(output.embedding, Nx.tensor([[4.0, 5.0, 6.0, 3.0]]))
+    end
+
+    @tag :tmp_dir
+    test "loads last_token pooling", %{spec: spec, tmp_dir: dir} do
+      custom_model =
+        Axon.input("input_ids", shape: {nil, nil})
+        |> Axon.nx(fn _input_ids ->
+          hidden_state =
+            Nx.tensor([[[1.0, 5.0, 2.0, 3.0], [4.0, 2.0, 6.0, 1.0], [99.0, 99.0, 99.0, 99.0]]])
+
+          %{hidden_state: hidden_state}
+        end)
+
+      model_info = %{model: custom_model, params: Axon.ModelState.empty(), spec: spec}
+
+      modules = [
+        %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"},
+        %{"idx" => 1, "path" => "1_Pooling", "type" => "sentence_transformers.models.Pooling"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      pooling_dir = Path.join(dir, "1_Pooling")
+      File.mkdir_p!(pooling_dir)
+
+      pooling_config = %{
+        "pooling_mode_cls_token" => false,
+        "pooling_mode_mean_tokens" => false,
+        "pooling_mode_lasttoken" => true
+      }
+
+      File.write!(Path.join(pooling_dir, "config.json"), Jason.encode!(pooling_config))
+
+      assert {:ok, new_model_info} = Bumblebee.load_embedding_head({:local, dir}, model_info)
+
+      inputs = %{
+        "input_ids" => Nx.tensor([[1, 2, 3]]),
+        "attention_mask" => Nx.tensor([[1, 1, 0]])
+      }
+
+      {_init_fn, predict_fn} = Axon.build(new_model_info.model)
+      output = predict_fn.(new_model_info.params, inputs)
+
+      assert_all_close(output.embedding, Nx.tensor([[4.0, 2.0, 6.0, 1.0]]))
+    end
+
+    @tag :tmp_dir
+    test "loads last_token pooling with pooling_mode_last_token key", %{spec: spec, tmp_dir: dir} do
+      custom_model =
+        Axon.input("input_ids", shape: {nil, nil})
+        |> Axon.nx(fn _input_ids ->
+          hidden_state =
+            Nx.tensor([[[1.0, 5.0, 2.0, 3.0], [4.0, 2.0, 6.0, 1.0], [99.0, 99.0, 99.0, 99.0]]])
+
+          %{hidden_state: hidden_state}
+        end)
+
+      model_info = %{model: custom_model, params: Axon.ModelState.empty(), spec: spec}
+
+      modules = [
+        %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"},
+        %{"idx" => 1, "path" => "1_Pooling", "type" => "sentence_transformers.models.Pooling"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      pooling_dir = Path.join(dir, "1_Pooling")
+      File.mkdir_p!(pooling_dir)
+
+      pooling_config = %{
+        "pooling_mode_cls_token" => false,
+        "pooling_mode_mean_tokens" => false,
+        "pooling_mode_last_token" => true
+      }
+
+      File.write!(Path.join(pooling_dir, "config.json"), Jason.encode!(pooling_config))
+
+      assert {:ok, new_model_info} = Bumblebee.load_embedding_head({:local, dir}, model_info)
+
+      inputs = %{
+        "input_ids" => Nx.tensor([[1, 2, 3]]),
+        "attention_mask" => Nx.tensor([[1, 1, 0]])
+      }
+
+      {_init_fn, predict_fn} = Axon.build(new_model_info.model)
+      output = predict_fn.(new_model_info.params, inputs)
+
+      assert_all_close(output.embedding, Nx.tensor([[4.0, 2.0, 6.0, 1.0]]))
+    end
+
+    @tag :tmp_dir
+    test "concatenates multiple pooling modes in PyTorch order (max before mean)", %{
+      spec: spec,
+      tmp_dir: dir
+    } do
+      custom_model =
+        Axon.input("input_ids", shape: {nil, nil})
+        |> Axon.nx(fn _input_ids ->
+          hidden_state =
+            Nx.tensor([[[1.0, 5.0, 2.0, 3.0], [4.0, 2.0, 6.0, 1.0], [99.0, 99.0, 99.0, 99.0]]])
+
+          %{hidden_state: hidden_state}
+        end)
+
+      model_info = %{model: custom_model, params: Axon.ModelState.empty(), spec: spec}
+
+      modules = [
+        %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"},
+        %{"idx" => 1, "path" => "1_Pooling", "type" => "sentence_transformers.models.Pooling"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      pooling_dir = Path.join(dir, "1_Pooling")
+      File.mkdir_p!(pooling_dir)
+
+      pooling_config = %{
+        "pooling_mode_cls_token" => false,
+        "pooling_mode_mean_tokens" => true,
+        "pooling_mode_max_tokens" => true
+      }
+
+      File.write!(Path.join(pooling_dir, "config.json"), Jason.encode!(pooling_config))
+
+      assert {:ok, new_model_info} = Bumblebee.load_embedding_head({:local, dir}, model_info)
+
+      inputs = %{
+        "input_ids" => Nx.tensor([[1, 2, 3]]),
+        "attention_mask" => Nx.tensor([[1, 1, 0]])
+      }
+
+      {_init_fn, predict_fn} = Axon.build(new_model_info.model)
+      output = predict_fn.(new_model_info.params, inputs)
+
+      # Max tokens: [4.0, 5.0, 6.0, 3.0]
+      # Mean tokens: [2.5, 3.5, 4.0, 2.0]
+      # Concatenated along last axis in PyTorch order (max then mean)
+      expected = Nx.tensor([[4.0, 5.0, 6.0, 3.0, 2.5, 3.5, 4.0, 2.0]])
+      assert_all_close(output.embedding, expected)
+    end
+
+    @tag :tmp_dir
+    test "preserves f16 type through pooling layer", %{spec: spec, tmp_dir: dir} do
+      custom_model =
+        Axon.input("input_ids", shape: {nil, nil})
+        |> Axon.nx(fn input_ids ->
+          batch_size = Nx.axis_size(input_ids, 0)
+          seq_len = Nx.axis_size(input_ids, 1)
+
+          hidden_state =
+            Nx.broadcast(1.0, {batch_size, seq_len, 4})
+            |> Nx.as_type({:f, 16})
+
+          %{hidden_state: hidden_state}
+        end)
+
+      model_info = %{model: custom_model, params: Axon.ModelState.empty(), spec: spec}
+
+      modules = [
+        %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"},
+        %{"idx" => 1, "path" => "1_Pooling", "type" => "sentence_transformers.models.Pooling"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      pooling_dir = Path.join(dir, "1_Pooling")
+      File.mkdir_p!(pooling_dir)
+
+      pooling_config = %{
+        "pooling_mode_cls_token" => false,
+        "pooling_mode_mean_tokens" => true,
+        "pooling_mode_max_tokens" => false
+      }
+
+      File.write!(Path.join(pooling_dir, "config.json"), Jason.encode!(pooling_config))
+
+      assert {:ok, new_model_info} = Bumblebee.load_embedding_head({:local, dir}, model_info)
+
+      inputs = %{
+        "input_ids" => Nx.tensor([[1, 2, 3]]),
+        "attention_mask" => Nx.tensor([[1, 1, 0]], type: {:u, 32})
+      }
+
+      {_init_fn, predict_fn} = Axon.build(new_model_info.model)
+      output = predict_fn.(new_model_info.params, inputs)
+
+      assert Nx.type(output.embedding) == {:f, 16}
+    end
+
+    @tag :tmp_dir
+    test "loads mean_sqrt_len_tokens pooling", %{spec: spec, tmp_dir: dir} do
+      custom_model =
+        Axon.input("input_ids", shape: {nil, nil})
+        |> Axon.nx(fn _input_ids ->
+          hidden_state =
+            Nx.tensor([[[1.0, 5.0, 2.0, 3.0], [4.0, 2.0, 6.0, 1.0], [99.0, 99.0, 99.0, 99.0]]])
+
+          %{hidden_state: hidden_state}
+        end)
+
+      model_info = %{model: custom_model, params: Axon.ModelState.empty(), spec: spec}
+
+      modules = [
+        %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"},
+        %{"idx" => 1, "path" => "1_Pooling", "type" => "sentence_transformers.models.Pooling"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      pooling_dir = Path.join(dir, "1_Pooling")
+      File.mkdir_p!(pooling_dir)
+
+      pooling_config = %{
+        "pooling_mode_cls_token" => false,
+        "pooling_mode_mean_tokens" => false,
+        "pooling_mode_mean_sqrt_len_tokens" => true
+      }
+
+      File.write!(Path.join(pooling_dir, "config.json"), Jason.encode!(pooling_config))
+
+      assert {:ok, new_model_info} = Bumblebee.load_embedding_head({:local, dir}, model_info)
+
+      inputs = %{
+        "input_ids" => Nx.tensor([[1, 2, 3]]),
+        "attention_mask" => Nx.tensor([[1, 1, 0]])
+      }
+
+      {_init_fn, predict_fn} = Axon.build(new_model_info.model)
+      output = predict_fn.(new_model_info.params, inputs)
+
+      expected = Nx.divide(Nx.tensor([[5.0, 7.0, 8.0, 4.0]]), :math.sqrt(2))
+      assert_all_close(output.embedding, expected)
+    end
+
+    @tag :tmp_dir
+    test "returns error when no supported pooling mode is enabled", %{
+      model_info: model_info,
+      tmp_dir: dir
+    } do
+      modules = [
+        %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"},
+        %{"idx" => 1, "path" => "1_Pooling", "type" => "sentence_transformers.models.Pooling"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      pooling_dir = Path.join(dir, "1_Pooling")
+      File.mkdir_p!(pooling_dir)
+
+      pooling_config = %{
+        "pooling_mode_cls_token" => false,
+        "pooling_mode_mean_tokens" => false
+      }
+
+      File.write!(Path.join(pooling_dir, "config.json"), Jason.encode!(pooling_config))
+
+      assert {:error, "no supported pooling mode found in 1_Pooling"} =
+               Bumblebee.load_embedding_head({:local, dir}, model_info)
+    end
+
+    @tag :tmp_dir
+    test "loads dense layer with activation (Tanh and ReLU)", %{
+      model_info: model_info,
+      tmp_dir: dir
+    } do
+      modules = [
+        %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"},
+        %{"idx" => 1, "path" => "1_Pooling", "type" => "sentence_transformers.models.Pooling"},
+        %{"idx" => 2, "path" => "2_Dense", "type" => "sentence_transformers.models.Dense"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      pooling_dir = Path.join(dir, "1_Pooling")
+      File.mkdir_p!(pooling_dir)
+
+      File.write!(
+        Path.join(pooling_dir, "config.json"),
+        Jason.encode!(%{"pooling_mode_mean_tokens" => true})
+      )
+
+      dense_dir = Path.join(dir, "2_Dense")
+      File.mkdir_p!(dense_dir)
+
+      # 1. Tanh
+      File.write!(
+        Path.join(dense_dir, "config.json"),
+        Jason.encode!(%{
+          "in_features" => 4,
+          "out_features" => 4,
+          "bias" => false,
+          "activation_function" => "torch.nn.modules.activation.Tanh"
+        })
+      )
+
+      weights = %{
+        "linear.weight" => Nx.broadcast(0.5, {4, 4}) |> Nx.as_type({:f, 32})
+      }
+
+      Safetensors.write!(Path.join(dense_dir, "model.safetensors"), weights)
+
+      assert {:ok, new_model_info} = Bumblebee.load_embedding_head({:local, dir}, model_info)
+
+      inputs = %{
+        "input_ids" => Nx.tensor([[1, 2, 3]]),
+        "attention_mask" => Nx.tensor([[1, 1, 1]])
+      }
+
+      {_init_fn, predict_fn} = Axon.build(new_model_info.model)
+      output = predict_fn.(new_model_info.params, inputs)
+
+      expected = Nx.tensor([[0.96402758, 0.96402758, 0.96402758, 0.96402758]])
+      assert_all_close(output.embedding, expected)
+
+      # 2. ReLU with negative weights
+      File.write!(
+        Path.join(dense_dir, "config.json"),
+        Jason.encode!(%{
+          "in_features" => 4,
+          "out_features" => 4,
+          "bias" => false,
+          "activation_function" => "torch.nn.modules.activation.ReLU"
+        })
+      )
+
+      neg_weights = %{
+        "linear.weight" => Nx.broadcast(-0.5, {4, 4}) |> Nx.as_type({:f, 32})
+      }
+
+      Safetensors.write!(Path.join(dense_dir, "model.safetensors"), neg_weights)
+
+      assert {:ok, relu_model_info} = Bumblebee.load_embedding_head({:local, dir}, model_info)
+
+      {_init_fn, predict_relu} = Axon.build(relu_model_info.model)
+      output_relu = predict_relu.(relu_model_info.params, inputs)
+
+      assert_all_close(output_relu.embedding, Nx.broadcast(0.0, {1, 4}))
+    end
+
+    @tag :tmp_dir
+    test "returns error on unsupported activation function in dense config", %{
+      model_info: model_info,
+      tmp_dir: dir
+    } do
+      modules = [
+        %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"},
+        %{"idx" => 1, "path" => "1_Pooling", "type" => "sentence_transformers.models.Pooling"},
+        %{"idx" => 2, "path" => "2_Dense", "type" => "sentence_transformers.models.Dense"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      pooling_dir = Path.join(dir, "1_Pooling")
+      File.mkdir_p!(pooling_dir)
+
+      File.write!(
+        Path.join(pooling_dir, "config.json"),
+        Jason.encode!(%{"pooling_mode_mean_tokens" => true})
+      )
+
+      dense_dir = Path.join(dir, "2_Dense")
+      File.mkdir_p!(dense_dir)
+
+      File.write!(
+        Path.join(dense_dir, "config.json"),
+        Jason.encode!(%{
+          "in_features" => 4,
+          "out_features" => 4,
+          "bias" => false,
+          "activation_function" => "torch.nn.modules.activation.Sigmoid"
+        })
+      )
+
+      assert {:error,
+              "unsupported activation function \"torch.nn.modules.activation.Sigmoid\" in 2_Dense"} =
+               Bumblebee.load_embedding_head({:local, dir}, model_info)
+    end
+
+    @tag :tmp_dir
+    test "handles Dropout module as pass-through", %{model_info: model_info, tmp_dir: dir} do
+      modules = [
+        %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"},
+        %{"idx" => 1, "path" => "1_Pooling", "type" => "sentence_transformers.models.Pooling"},
+        %{"idx" => 2, "path" => "2_Dropout", "type" => "sentence_transformers.models.Dropout"},
+        %{"idx" => 3, "path" => "3_Normalize", "type" => "sentence_transformers.models.Normalize"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      pooling_dir = Path.join(dir, "1_Pooling")
+      File.mkdir_p!(pooling_dir)
+
+      File.write!(
+        Path.join(pooling_dir, "config.json"),
+        Jason.encode!(%{"pooling_mode_mean_tokens" => true})
+      )
+
+      assert {:ok, new_model_info} = Bumblebee.load_embedding_head({:local, dir}, model_info)
+
+      inputs = %{
+        "input_ids" => Nx.tensor([[1, 2, 3]]),
+        "attention_mask" => Nx.tensor([[1, 1, 1]])
+      }
+
+      {_init_fn, predict_fn} = Axon.build(new_model_info.model)
+      output = predict_fn.(new_model_info.params, inputs)
+
+      assert Nx.shape(output.embedding) == {1, 4}
+      norm = Nx.LinAlg.norm(output.embedding, axes: [-1])
+      assert_all_close(norm, Nx.tensor([1.0]), atol: 1.0e-5)
+    end
+
+    @tag :tmp_dir
+    test "loads dense layer parameters from pytorch_model.bin", %{
+      model_info: model_info,
+      tmp_dir: dir
+    } do
+      modules = [
+        %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"},
+        %{"idx" => 1, "path" => "1_Pooling", "type" => "sentence_transformers.models.Pooling"},
+        %{"idx" => 2, "path" => "2_Dense", "type" => "sentence_transformers.models.Dense"}
+      ]
+
+      File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules))
+
+      pooling_dir = Path.join(dir, "1_Pooling")
+      File.mkdir_p!(pooling_dir)
+
+      File.write!(
+        Path.join(pooling_dir, "config.json"),
+        Jason.encode!(%{"pooling_mode_mean_tokens" => true})
+      )
+
+      dense_dir = Path.join(dir, "2_Dense")
+      File.mkdir_p!(dense_dir)
+
+      File.write!(
+        Path.join(dense_dir, "config.json"),
+        Jason.encode!(%{
+          "in_features" => 4,
+          "out_features" => 8,
+          "bias" => true,
+          "activation_function" => "torch.nn.modules.linear.Identity"
+        })
+      )
+
+      pkl =
+        <<128, 2, "ccollections\nOrderedDict\nq", 0, ")Rq", 1, "(X", 13, 0, 0, 0,
+          "linear.weightq", 2, "ctorch._utils\n_rebuild_tensor_v2\nq", 3, "((X", 7, 0, 0, 0,
+          "storageq", 4, "ctorch\nFloatStorage\nq", 5, "X", 1, 0, 0, 0, "0q", 6, "X", 3, 0, 0, 0,
+          "cpuq", 7, "K ", 116, 113, 8, "QK", 0, "(K", 8, "K", 4, 116, 113, 9, "(K", 4, "K", 1,
+          116, 113, 10, 137, 104, 0, ")Rq", 11, 116, 113, 12, "Rq", 13, "X", 11, 0, 0, 0,
+          "linear.biasq", 14, "h", 3, "((h", 4, "h", 5, "X", 1, 0, 0, 0, "1q", 15, "h", 7, "K", 8,
+          116, 113, 16, "QK", 0, "K", 8, 133, 113, 17, "K", 1, 133, 113, 18, 137, 104, 0, ")Rq",
+          19, 116, 113, 20, "Rq", 21, "u.">>
+
+      data0 = :binary.copy(<<0, 0, 0, 63>>, 32)
+      data1 = :binary.copy(<<205, 204, 204, 61>>, 8)
+
+      files = [
+        {~c"archive/data.pkl", pkl},
+        {~c"archive/data/0", data0},
+        {~c"archive/data/1", data1},
+        {~c"archive/version", "3\n"}
+      ]
+
+      pt_path = Path.join(dense_dir, "pytorch_model.bin")
+      {:ok, _} = :zip.create(to_charlist(pt_path), files)
+
+      assert {:ok, new_model_info} = Bumblebee.load_embedding_head({:local, dir}, model_info)
+
+      assert Map.has_key?(new_model_info.params.data, "2_Dense")
+      assert Nx.shape(new_model_info.params.data["2_Dense"]["kernel"]) == {4, 8}
+      assert Nx.shape(new_model_info.params.data["2_Dense"]["bias"]) == {8}
+
+      inputs = %{
+        "input_ids" => Nx.tensor([[1, 2, 3]]),
+        "attention_mask" => Nx.tensor([[1, 1, 1]])
+      }
+
+      {_init_fn, predict_fn} = Axon.build(new_model_info.model)
+      output = predict_fn.(new_model_info.params, inputs)
+
+      assert Nx.shape(output.embedding) == {1, 8}
+    end
+  end
+
+  @tag :slow
+  test "end-to-end with unsloth/embeddinggemma-300m matches reference SentenceTransformers output" do
+    assert {:ok, base_model} = Bumblebee.load_model({:hf, "unsloth/embeddinggemma-300m"})
+    assert {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "unsloth/embeddinggemma-300m"})
+
+    # 1. Standard Separate Head
+    assert {:ok, model_info} =
+             Bumblebee.load_embedding_head({:hf, "unsloth/embeddinggemma-300m"}, base_model)
+
+    serving = Bumblebee.Text.text_embedding(model_info, tokenizer)
+    res = Nx.Serving.run(serving, "Hello world")
+
+    assert Nx.shape(res.embedding) == {768}
+
+    assert_all_close(
+      res.embedding[0..4],
+      Nx.tensor([-0.203079, 0.034759, 0.060166, -0.016863, 0.006666]),
+      atol: 1.0e-4
+    )
+
+    # 2. Optimized Fused Head
+    assert {:ok, fused_model_info} =
+             Bumblebee.load_embedding_head(
+               {:hf, "unsloth/embeddinggemma-300m"},
+               base_model,
+               fuse_dense: true
+             )
+
+    fused_serving = Bumblebee.Text.text_embedding(fused_model_info, tokenizer)
+    fused_res = Nx.Serving.run(fused_serving, "Hello world")
+
+    assert_all_close(res.embedding, fused_res.embedding)
+  end
+
+  @tag :slow
+  test "end-to-end with sentence-transformers/all-MiniLM-L6-v2" do
+    assert {:ok, model_info} =
+             Bumblebee.load_model({:hf, "sentence-transformers/all-MiniLM-L6-v2"})
+
+    assert {:ok, model_info} =
+             Bumblebee.load_embedding_head(
+               {:hf, "sentence-transformers/all-MiniLM-L6-v2"},
+               model_info
+             )
+
+    assert {:ok, tokenizer} =
+             Bumblebee.load_tokenizer({:hf, "sentence-transformers/all-MiniLM-L6-v2"})
+
+    serving = Bumblebee.Text.text_embedding(model_info, tokenizer)
+    res = Nx.Serving.run(serving, "Hello world")
+
+    assert Nx.shape(res.embedding) == {384}
+
+    norm = Nx.LinAlg.norm(res.embedding)
+    assert_all_close(norm, Nx.tensor(1.0), atol: 1.0e-5)
+  end
+
+  @tag :slow
+  test "end-to-end with sentence-transformers/distiluse-base-multilingual-cased-v1 (Dense with Tanh)" do
+    assert {:ok, model_info} =
+             Bumblebee.load_model(
+               {:hf, "sentence-transformers/distiluse-base-multilingual-cased-v1"}
+             )
+
+    assert {:ok, model_info} =
+             Bumblebee.load_embedding_head(
+               {:hf, "sentence-transformers/distiluse-base-multilingual-cased-v1"},
+               model_info
+             )
+
+    assert {:ok, tokenizer} =
+             Bumblebee.load_tokenizer(
+               {:hf, "sentence-transformers/distiluse-base-multilingual-cased-v1"}
+             )
+
+    serving = Bumblebee.Text.text_embedding(model_info, tokenizer)
+    res = Nx.Serving.run(serving, "Hello world")
+
+    assert Nx.shape(res.embedding) == {512}
+
+    assert_all_close(
+      res.embedding[0..4],
+      Nx.tensor([0.037437, 0.000251, -0.044534, -0.006018, -0.004641]),
+      atol: 1.0e-4
+    )
+  end
+
+  @tag :slow
+  test "end-to-end with sentence-transformers/sentence-t5-base (T5 encoder with Dense projection)" do
+    assert {:ok, model_info} =
+             Bumblebee.load_model({:hf, "sentence-transformers/sentence-t5-base"})
+
+    assert {:ok, model_info} =
+             Bumblebee.load_embedding_head(
+               {:hf, "sentence-transformers/sentence-t5-base"},
+               model_info
+             )
+
+    assert {:ok, tokenizer} =
+             Bumblebee.load_tokenizer({:hf, "sentence-transformers/sentence-t5-base"})
+
+    serving = Bumblebee.Text.text_embedding(model_info, tokenizer)
+    res = Nx.Serving.run(serving, "Hello world")
+
+    assert Nx.shape(res.embedding) == {768}
+
+    norm = Nx.LinAlg.norm(res.embedding)
+    assert_all_close(norm, Nx.tensor(1.0), atol: 1.0e-5)
+
+    assert_all_close(
+      res.embedding[0..4],
+      Nx.tensor([0.001551, -0.056042, 0.026304, 0.058385, 0.007024]),
+      atol: 1.0e-4
+    )
+  end
+
+  @tag :slow
+  test "end-to-end with sentence-transformers/all-distilroberta-v1 (architecture: :base)" do
+    assert {:ok, model_info} =
+             Bumblebee.load_model(
+               {:hf, "sentence-transformers/all-distilroberta-v1"},
+               architecture: :base
+             )
+
+    assert {:ok, model_info} =
+             Bumblebee.load_embedding_head(
+               {:hf, "sentence-transformers/all-distilroberta-v1"},
+               model_info
+             )
+
+    assert {:ok, tokenizer} =
+             Bumblebee.load_tokenizer({:hf, "sentence-transformers/all-distilroberta-v1"})
+
+    serving = Bumblebee.Text.text_embedding(model_info, tokenizer)
+    res = Nx.Serving.run(serving, "Hello world")
+
+    assert Nx.shape(res.embedding) == {768}
+
+    norm = Nx.LinAlg.norm(res.embedding)
+    assert_all_close(norm, Nx.tensor(1.0), atol: 1.0e-5)
+
+    assert_all_close(
+      res.embedding[0..4],
+      Nx.tensor([0.027808, -0.024371, -0.021978, -0.057262, 0.034520]),
+      atol: 1.0e-4
+    )
+  end
+end