Add Bumblebee.load_embedding_head for SentenceTransformers pipelines

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}
 

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 """