[dkuku]

Currently everything is self contained in the new module Bumblebee.SentenceTransformers.load_model(repo), but it can be exposed in Bumblebee.

I had more tests that compare output with python - but it is removed in one of the commits. It can be restored or used for testing this pr.

Assisted with codex and antigravity

Summary

Adds support for SentenceTransformers models via Bumblebee.SentenceTransformers.

SentenceTransformers repositories define multi-stage inference pipelines in modules.json that stack pooling, dense projections, normalization, and custom prompt logic on top of base transformer or static embedding models. This PR adds dynamic loading and serving for these architectures while integrating natively with Bumblebee's existing Nx.Serving ecosystem.

---

Key Additions

• Bumblebee.SentenceTransformers:

  • load_model/2 & load_tokenizer/2: Resolves and loads the base model and stacks pipeline modules from modules.json (with automatic fallback pooling if modules.json is missing).
  • load_embedding_head/3: Attaches downstream SentenceTransformers layers onto an existing model_info.
  • text_embedding/3: Builds an Nx.Serving pipeline handling prompts (from config_sentence_transformers.json or runtime overrides), prompt exclusion from pooling, and max_seq_length truncation.
  • encode_queries/4 & encode_corpus/4: Convenience helpers resolving default prompt aliases ("query", "document", etc.).
  • truncate_embeddings/3: Matryoshka dimension truncation with optional L2 re-normalization.

• Pipeline Modules Support:

  • Pooling (mean, cls, max, weightedmean, lasttoken, mean_sqrt_len_tokens)
  • Dense projections (linear layers, activations, residuals, and optional dense fusing)
  • Normalization (Normalize, LayerNorm)
  • WeightedLayerPooling & StaticEmbedding (e.g. Model2Vec)

---

Usage

  # 1. Load model and tokenizer
  repo = {:hf, "sentence-transformers/all-MiniLM-L6-v2"}
  {:ok, model_info} = Bumblebee.SentenceTransformers.load_model(repo)
  {:ok, tokenizer} = Bumblebee.SentenceTransformers.load_tokenizer(repo)
  
  # 2. Build serving and run embedding
  serving = Bumblebee.SentenceTransformers.text_embedding(model_info, tokenizer)
  Nx.Serving.run(serving, "Hello world")
  #=> %{embedding: #Nx.Tensor<f32[384] ...>}
  
  # 3. Asymmetric search with prompt templates queries = Bumblebee.SentenceTransformers.encode_queries(model_info, tokenizer, ["Where is Paris?"])
    corpus = Bumblebee.SentenceTransformers.encode_corpus(model_info, tokenizer, ["Paris is in France."])

  # 4. Attach head to an existing base model (optional)
  {:ok, model_info} =
    Bumblebee.load_model(repo, architecture: :base)
    |> Bumblebee.SentenceTransformers.load_embedding_head(repo)

[g-andrade]

FWIW, EmbeddingGemma also needs #480 for inputs over 256 tokens.