Testing models Elixir vs Python: compute embeddings with "sentence-transformers/all-MiniLM-L6-v2"

[ndrean]

We will check that Elixir returns the same values as Python (considered as the source of truth).

Python code: jupyter notebook

from sentence_transformers import SentenceTransformer
import numpy as np

def load_model(model_name="sentence-transformers/all-MiniLM-L6-v2"):
    return SentenceTransformer(model_name)

def generate_embedding(model, text):
    # Clean the text (similar to String.trim in Elixir)
    text = text.strip()
    # Compute the embedding with L2 normalization
    embedding = model.encode(text, normalize_embeddings=True)
    return embedding

model = load_model()

text = "This is an example sentence for embedding."
embedding = generate_embedding(model, text)

print("Embedding:", embedding)

Elixiir code: livebook server

Mix.install(
[{:bumblebee, "~> 0.5.3"},{:exla, "~> 0.7.3"}, {:kino, "~> 0.14.0"},], 
 config: [nx: [default_backend: EXLA.Backend]])

Nx.Defn.global_default_options(compiler: EXLA, client: :host)

defmodule Embedder do

  def load_model do
    repo = {:hf, "sentence-transformers/all-MiniLM-L6-v2"}
    {:ok, model_info} = Bumblebee.load_model(repo) 
    {:ok, tokenizer} = Bumblebee.load_tokenizer(repo)

    embedding_serving = 
      Bumblebee.Text.text_embedding(
        model_info, 
        tokenizer,
        output_pool: :mean_pooling,
        output_attribute: :hidden_state,
        embedding_processor: :l2_norm,
        compile: [batch_size: 1, sequence_length: [2000]],
        defn_options: [compiler: EXLA]
      )

    Kino.start_child({Nx.Serving, serving: embedding_serving, name: ChunkEmbedder})
  end

  def generate_embedding(text) do
    %{embedding: vector} = Nx.Serving.batched_run(ChunkEmbedder, String.trim(text))
    Nx.to_flat_list(vector)
  end
end

Embedder.load_model()

Embedder.generate_embedding("This is an example sentence for embedding.")

Results

[0.043010398745536804, 0.009119901806116104, 0.05472865328192711, 0.04987872764468193,
 0.04459322616457939, 0.051743000745773315, 0.019657032564282417, -0.021903224289417267,
 0.04263623431324959, -0.040632933378219604, 0.056533899158239365,.....]