From b76f27a0a126b14cf5ded4ad38cd6a378acac2d3 Mon Sep 17 00:00:00 2001 From: Daniel Kukula Date: Tue, 6 Oct 2026 21:58:05 +0200 Subject: [PATCH 1/4] Add SentenceTransformers module pipeline and inference support --- lib/bumblebee/sentence_transformers.ex | 870 +++++++++++ .../sentence_transformers/downloader.ex | 108 ++ .../sentence_transformers/pipeline.ex | 816 +++++++++++ mix.exs | 4 + .../sentence_transformers/dense_edge_test.exs | 260 ++++ .../python_comparison_test.exs | 557 +++++++ .../python_environment_test.exs | 48 + test/bumblebee/sentence_transformers_test.exs | 1289 +++++++++++++++++ .../sentence_transformers/python_bridge.ex | 557 +++++++ 9 files changed, 4509 insertions(+) create mode 100644 lib/bumblebee/sentence_transformers.ex create mode 100644 lib/bumblebee/sentence_transformers/downloader.ex create mode 100644 lib/bumblebee/sentence_transformers/pipeline.ex create mode 100644 test/bumblebee/sentence_transformers/dense_edge_test.exs create mode 100644 test/bumblebee/sentence_transformers/python_comparison_test.exs create mode 100644 test/bumblebee/sentence_transformers/python_environment_test.exs create mode 100644 test/bumblebee/sentence_transformers_test.exs create mode 100644 test/support/sentence_transformers/python_bridge.ex diff --git a/lib/bumblebee/sentence_transformers.ex b/lib/bumblebee/sentence_transformers.ex new file mode 100644 index 00000000..431f8b9e --- /dev/null +++ b/lib/bumblebee/sentence_transformers.ex @@ -0,0 +1,870 @@ +defmodule Bumblebee.SentenceTransformers do + @moduledoc """ + SentenceTransformers pipelines in Elixir on top of Bumblebee, Axon, and Nx. + + SentenceTransformers supports a subset of Python's `sentence-transformers` inference modules + by parsing `modules.json` and downloading component files on demand + (e.g. `1_Pooling/config.json`, `2_Dense/model.safetensors`). + + ## Quick start + + {:ok, model_info} = Bumblebee.SentenceTransformers.load_model({:hf, "sentence-transformers/all-MiniLM-L6-v2"}) + {:ok, tokenizer} = Bumblebee.SentenceTransformers.load_tokenizer({:hf, "sentence-transformers/all-MiniLM-L6-v2"}) + + serving = Bumblebee.SentenceTransformers.text_embedding(model_info, tokenizer) + Nx.Serving.run(serving, "Hello world") + #=> %{embedding: #Nx.Tensor} + + `text_embedding/3` uses a saved default prompt and saved maximum sequence + length when the caller does not override them. + """ + + alias Bumblebee.SentenceTransformers.Downloader + alias Bumblebee.SentenceTransformers.Pipeline + alias Bumblebee.Text.PreTrainedTokenizer + + @type repository :: Bumblebee.repository() + + @doc """ + Loads a SentenceTransformers model from a repository. + + If `modules.json` is present in the repository, SentenceTransformers parses the module pipeline, + loads the underlying base model, and stacks the declared modules (pooling, dense + projections, normalization, etc.) on top. + + If `modules.json` is not present, it applies fallback pooling: `:last` for + causal language models and `:mean` for other model families. + + ## Options + + * `:fuse_dense` - whether to pre-multiply consecutive linear projection layers + without bias into a single projection layer. Defaults to `false` + + * `:type` - the numerical precision or Axon mixed precision policy + + * `:backend` - the Nx backend to allocate parameters on + + * `:safetensors_reader` - function used to read `.safetensors` files + + """ + @spec load_model(repository(), keyword()) :: {:ok, map()} | {:error, String.t()} + def load_model(repository, opts \\ []) do + base_opts = + opts + |> Keyword.take([:backend, :type, :safetensors_reader]) + |> Keyword.put_new(:architecture, :base) + + case Downloader.download_file(repository, "modules.json") do + {:ok, modules_path} -> + with {:ok, modules} <- Pipeline.decode_json(modules_path), + modules = Enum.sort_by(modules, &(&1["idx"] || 0)), + [first | pipeline_modules] <- modules do + first_type = Pipeline.normalize_module_type(first["type"]) + + case first_type do + "Transformer" -> + base_repository = Downloader.subrepository(repository, first["path"]) + + with {:ok, model_info} <- Bumblebee.load_model(base_repository, base_opts), + {:ok, model_info} <- + Pipeline.build(repository, pipeline_modules, model_info, opts) do + load_config_and_attach(repository, model_info) + end + + "StaticEmbedding" -> + with {:ok, model_info} <- + Pipeline.build_static_embedding(repository, first, opts), + {:ok, model_info} <- + Pipeline.build(repository, pipeline_modules, model_info, opts) do + load_config_and_attach(repository, model_info) + end + + other -> + {:error, "unsupported base module in modules.json: #{inspect(other)}"} + end + else + [] -> + {:error, "modules.json is empty"} + + {:error, reason} -> + {:error, reason} + + _other -> + {:error, + "expected modules.json to define a Transformer or StaticEmbedding base model"} + end + + {:error, reason} -> + if Downloader.missing_file?(repository, "modules.json", reason) do + load_fallback_model(repository, base_opts, opts) + else + {:error, reason} + end + end + end + + defp load_fallback_model(repository, base_opts, opts) do + # Fallback when modules.json is not present: load model and apply default pooling + # CausalLM-based models use last token pooling, otherwise mean pooling + with {:ok, model_info} <- Bumblebee.load_model(repository, base_opts) do + pooling_mode = infer_fallback_pooling_mode(repository, model_info) + + with {:ok, model_info} <- + Pipeline.build( + repository, + [ + %{ + "type" => "Pooling", + "path" => "default_pooling", + "pooling_mode" => pooling_mode + } + ], + model_info, + opts + ) do + load_config_and_attach(repository, model_info) + end + end + end + + @doc """ + Loads only the SentenceTransformers embedding head and attaches it to an existing `model_info`. + + Supports piping directly from `Bumblebee.load_model/2`: + + {:ok, model_info} = + Bumblebee.load_model(repo, architecture: :base) + |> Bumblebee.SentenceTransformers.load_embedding_head(repo) + + """ + def load_embedding_head(model_or_repo, repo_or_model, opts \\ []) + + def load_embedding_head({:ok, model_info}, repository, opts) when is_map(model_info) do + load_embedding_head(model_info, repository, opts) + end + + def load_embedding_head({:error, reason}, _repository, _opts) do + {:error, reason} + end + + def load_embedding_head(%{model: _} = model_info, repository, opts) do + do_load_embedding_head(repository, model_info, opts) + end + + def load_embedding_head(repository, %{model: _} = model_info, opts) do + do_load_embedding_head(repository, model_info, opts) + end + + def load_embedding_head(repository, {:ok, %{model: _} = model_info}, opts) do + do_load_embedding_head(repository, model_info, opts) + end + + @doc """ + Alias for `load_embedding_head/3`. + """ + def load_head(model_or_repo, repo_or_model, opts \\ []) do + load_embedding_head(model_or_repo, repo_or_model, opts) + end + + defp do_load_embedding_head(repository, model_info, opts) do + case Downloader.download_file(repository, "modules.json") do + {:ok, modules_path} -> + with {:ok, modules} <- Pipeline.decode_json(modules_path), + modules = Enum.sort_by(modules, &(&1["idx"] || 0)), + [first | pipeline_modules] <- modules, + true <- + Pipeline.normalize_module_type(first["type"]) in ["Transformer", "StaticEmbedding"] || + {:error, + "expected the first module in modules.json to be Transformer or StaticEmbedding"}, + {:ok, model_info} <- Pipeline.build(repository, pipeline_modules, model_info, opts) do + load_config_and_attach(repository, model_info) + else + {:error, reason} -> + {:error, reason} + + _other -> + {:error, + "expected modules.json to define a Transformer or StaticEmbedding base model"} + end + + {:error, _} -> + {:error, "could not find modules.json in the repository"} + end + end + + @doc """ + Loads the tokenizer for the first module in a SentenceTransformers pipeline, + or the repository root when `modules.json` is absent. + """ + def load_tokenizer(repository, opts \\ []) do + case Downloader.download_file(repository, "modules.json") do + {:ok, path} -> + with {:ok, modules} <- Pipeline.decode_json(path), + [first | _] <- Enum.sort_by(modules, &(&1["idx"] || 0)) do + Bumblebee.load_tokenizer(Downloader.subrepository(repository, first["path"]), opts) + else + {:error, reason} -> {:error, reason} + _ -> {:error, "modules.json is empty"} + end + + {:error, reason} -> + if Downloader.missing_file?(repository, "modules.json", reason) do + Bumblebee.load_tokenizer(repository, opts) + else + {:error, reason} + end + end + end + + @doc """ + Builds a text embedding serving that returns `%{embedding: tensor}` for one + input and a list of those maps for a list of inputs. Set + `return_tokenized: true` to also return each input's tokenized + `%{"input_ids" => tensor, "attention_mask" => tensor}`. + + This mirrors `Bumblebee.Text.text_embedding/3` and additionally supports + SentenceTransformers models configured with `include_prompt: false`. In those + cases the prompt prefix is tokenized separately and a `prompt_length` input is + passed to the model so that prompt tokens are excluded from pooling. + + ## Options + + * `:prompt` - a custom prompt string to prepend to every input. It takes + precedence over the saved default prompt. Pass `false` or `nil` to + disable prompt use. + + * `:max_seq_length` - overrides the saved token limit. Compilation bucket + sizes control padding; they do not increase this limit. Smaller buckets + cap it at the largest available bucket. + + * `:prompt_name` - a saved prompt name from + `config_sentence_transformers.json`. It cannot be combined with + `:prompt`; pass `false` to disable prompt use. With neither + option, the saved default prompt is used. + + * `:output_attribute` - model output key to embed. Defaults to + `:pooled_state`. + + * `:output_pool` - optional `:cls_token_pooling`, `:mean_pooling`, or + `:last_token_pooling` for rank-3 outputs. Defaults to `nil`. + + * `:embedding_processor` - optional `:l2_norm`. Defaults to `nil`. + + * `:return_tokenized` - includes tokenized inputs in each result map. + Defaults to `false`. + + * `:compile`, `:defn_options`, `:preallocate_params` - serving execution + options. A compile sequence length sets a padding bucket and caps the + effective saved or requested maximum sequence length. + + """ + def text_embedding(model_info, tokenizer, opts \\ []) do + tokenizer = + if model_info[:sparse_static_embedding] do + Bumblebee.configure(tokenizer, add_special_tokens: false) + else + tokenizer + end + + text_embedding_with_prompt_length(model_info, tokenizer, opts) + end + + @doc false + def configure_embedding_tokenizer(model_info, tokenizer, opts) do + saved = get_in(model_info, [:sentence_transformers, :max_seq_length]) + limit = Keyword.get(opts, :max_seq_length, saved) + + if limit != nil and (not is_integer(limit) or limit <= 0) do + raise ArgumentError, ":max_seq_length must be a positive integer" + end + + buckets = get_in(opts, [:compile, :sequence_length]) + upper = if buckets, do: Enum.max(List.wrap(buckets)) + limit = if limit && upper, do: min(limit, upper), else: limit || upper + tokenizer = Bumblebee.configure(tokenizer, length: nil, return_token_type_ids: false) + + if limit do + # Bumblebee's length option couples truncation with fixed padding. Set + # native truncation alone so short inputs still use dynamic/small buckets. + update_in(tokenizer.native_tokenizer, fn native -> + Tokenizers.Tokenizer.set_truncation(native, + max_length: limit, + direction: tokenizer.truncate_direction + ) + end) + else + tokenizer + end + end + + @doc false + def apply_embedding_tokenizer(tokenizer, texts, buckets) do + inputs = Bumblebee.apply_tokenizer(tokenizer, texts) + + if buckets do + length = Nx.axis_size(inputs["input_ids"], 1) + target = buckets |> List.wrap() |> Enum.sort() |> Enum.find(&(&1 >= length)) + + pad_id = + PreTrainedTokenizer.token_to_id( + tokenizer, + tokenizer.special_tokens[:pad] || tokenizer.special_tokens[:eos] + ) + + Map.new(inputs, fn {key, tensor} -> + value = if key == "input_ids", do: pad_id, else: 0 + + padding = + if tokenizer.pad_direction == :left, + do: {target - length, 0, 0}, + else: {0, target - length, 0} + + {key, Nx.pad(tensor, value, [{0, 0, 0}, padding])} + end) + else + inputs + end + end + + @doc """ + Builds a cross encoding serving using `Bumblebee.Text.cross_encoding/3`. + """ + defdelegate cross_encoding(model_info, tokenizer, opts \\ []), to: Bumblebee.Text + + @doc """ + Encodes query text using the saved `"query"` prompt when available. + + This is a convenience wrapper around `text_embedding/3`. If no `"query"` + prompt is configured, it falls back to the saved default prompt (if any). + + Explicit `:prompt` or `:prompt_name` options take precedence. Returns the + same `%{embedding: tensor}` result shape as `text_embedding/3`. + """ + def encode_queries(model_info, tokenizer, queries, opts \\ []) do + serving = + text_embedding(model_info, tokenizer, encoding_prompt_opts(model_info, opts, ["query"])) + + Nx.Serving.run(serving, queries) + end + + @doc """ + Encodes corpus/document text using the first available saved prompt named + `"document"`, `"passage"`, or `"corpus"`. + + If none of those aliases are configured, it falls back to the saved default + prompt (if any). + + Explicit `:prompt` or `:prompt_name` options take precedence. Returns the + same `%{embedding: tensor}` result shape as `text_embedding/3`. + """ + def encode_corpus(model_info, tokenizer, corpus, opts \\ []) do + serving = + text_embedding( + model_info, + tokenizer, + encoding_prompt_opts(model_info, opts, ["document", "passage", "corpus"]) + ) + + Nx.Serving.run(serving, corpus) + end + + defp encoding_prompt_opts(model_info, opts, names) do + prompts = get_in(model_info, [:sentence_transformers, :prompts]) || %{} + + if Keyword.has_key?(opts, :prompt) or Keyword.has_key?(opts, :prompt_name) do + opts + else + case Enum.find(names, &Map.has_key?(prompts, &1)) do + nil -> opts + name -> Keyword.put(opts, :prompt_name, name) + end + end + end + + defp text_embedding_with_prompt_length(model_info, tokenizer, opts) do + %{model: model, params: params, spec: _spec} = model_info + + opts = + Keyword.validate!(opts, [ + :compile, + :prompt, + :prompt_name, + :max_seq_length, + return_tokenized: false, + output_attribute: :pooled_state, + output_pool: nil, + embedding_processor: nil, + defn_options: [], + preallocate_params: false + ]) + + prompt_prefix = resolve_prompt(opts, model_info) + + output_attribute = opts[:output_attribute] + output_pool = opts[:output_pool] + embedding_processor = opts[:embedding_processor] + preallocate_params = opts[:preallocate_params] + defn_options = opts[:defn_options] + + compile = + if compile = opts[:compile] do + compile + |> Keyword.validate!([:batch_size, :sequence_length]) + |> Bumblebee.Shared.require_options!([:batch_size, :sequence_length]) + end + + batch_size = compile[:batch_size] + sequence_length = compile[:sequence_length] + + tokenizer = + configure_embedding_tokenizer(model_info, tokenizer, opts) + + prompt_length = prompt_token_length(tokenizer, prompt_prefix) + + {_init_fun, encoder} = Axon.build(model) + + embedding_fun = fn params, inputs -> + output = encoder.(params, inputs) + + output = + case output do + %Nx.Tensor{} -> + output + + %{^output_attribute => %Axon.None{}} -> + keys = output |> Map.keys() |> Enum.sort() + + raise ArgumentError, + "key #{inspect(output_attribute)} in the output map has value %Axon.None{}," <> + " make sure this is the correct key and check module documentation incase it is opt-in." <> + " The output map keys are: #{inspect(keys)}" + + %{^output_attribute => output} -> + output + + %{embedding: output} when output_attribute == :pooled_state -> + output + + %{} -> + keys = output |> Map.keys() |> Enum.sort() + + raise ArgumentError, + "key #{inspect(output_attribute)} not found in the output map," <> + " you may want to set :output_attribute to one of the map keys: #{inspect(keys)}" + + _ -> + output + end + + if output_pool != nil and Nx.rank(output) != 3 do + raise ArgumentError, + "expected the output tensor to have rank 3 to apply :output_pool, got: #{Nx.rank(output)}." <> + " You should either disable pooling or pick a different output using :output_attribute" + end + + output = + case output_pool do + nil -> + output + + :cls_token_pooling -> + Nx.take(output, 0, axis: 1) + + :mean_pooling -> + input_mask_expanded = Nx.new_axis(inputs["attention_mask"], -1) + + output + |> Nx.multiply(input_mask_expanded) + |> Nx.sum(axes: [1]) + |> Nx.divide(Nx.sum(input_mask_expanded, axes: [1])) + + :last_token_pooling -> + sequence_lengths = + inputs["attention_mask"] + |> Nx.sum(axes: [1]) + |> Nx.subtract(1) + |> Nx.as_type({:s, 64}) + + Bumblebee.Utils.Nx.batched_take(output, sequence_lengths) + + other -> + raise ArgumentError, + "expected :output_pool to be one of :cls_token_pooling, :mean_pooling, :last_token_pooling or nil, got: #{inspect(other)}" + end + + output = + case embedding_processor do + nil -> + output + + :l2_norm -> + Bumblebee.Utils.Nx.normalize(output) + + other -> + raise ArgumentError, + "expected :embedding_processor to be one of nil or :l2_norm, got: #{inspect(other)}" + end + + output + end + + batch_keys = Bumblebee.Shared.sequence_batch_keys(sequence_length) + + fn batch_key, defn_options -> + params = Bumblebee.Shared.maybe_preallocate(params, preallocate_params, defn_options) + + scope = {:embedding, batch_key} + + embedding_fun = + Bumblebee.Shared.compile_or_jit( + embedding_fun, + scope, + defn_options, + compile != nil, + fn -> + {:sequence_length, sequence_length} = batch_key + + inputs = %{ + "input_ids" => Nx.template({batch_size, sequence_length}, :u32), + "attention_mask" => Nx.template({batch_size, sequence_length}, :u32), + "prompt_length" => Nx.template({batch_size}, :u32) + } + + [params, inputs] + end + ) + + fn inputs -> + inputs = Bumblebee.Shared.maybe_pad(inputs, batch_size) + params |> embedding_fun.(inputs) |> Bumblebee.Shared.serving_post_computation() + end + end + |> Nx.Serving.new(defn_options) + |> Nx.Serving.batch_size(batch_size) + |> Nx.Serving.process_options(batch_keys: batch_keys) + |> Nx.Serving.client_preprocessing(fn input -> + {texts, multi?} = + Bumblebee.Shared.validate_serving_input!(input, &Bumblebee.Shared.validate_string/1) + + texts = + case prompt_prefix do + "" -> texts + nil -> texts + prefix -> Enum.map(texts, &(prefix <> &1)) + end + + inputs = + Nx.with_default_backend(Nx.BinaryBackend, fn -> + apply_embedding_tokenizer(tokenizer, texts, sequence_length) + end) + + batch_size = Nx.axis_size(inputs["input_ids"], 0) + inputs = Map.put(inputs, "prompt_length", Nx.broadcast(prompt_length, {batch_size})) + + batch_key = Bumblebee.Shared.sequence_batch_key_for_inputs(inputs, sequence_length) + batch = [inputs] |> Nx.Batch.concatenate() |> Nx.Batch.key(batch_key) + + tokenized = + if opts[:return_tokenized], do: Map.take(inputs, ["input_ids", "attention_mask"]) + + {batch, {multi?, tokenized}} + end) + |> Nx.Serving.client_postprocessing(fn {embeddings, _metadata}, {multi?, tokenized} -> + embeddings + |> Bumblebee.Utils.Nx.batch_to_list() + |> Enum.with_index() + |> Enum.map(fn {embedding, index} -> + if tokenized do + row = + Map.new(tokenized, fn {key, tensor} -> + {key, Nx.slice_along_axis(tensor, index, 1, axis: 0)} + end) + + %{embedding: embedding, tokenized: row} + else + %{embedding: embedding} + end + end) + |> Bumblebee.Shared.normalize_output(multi?) + end) + end + + defp prompt_token_length(_tokenizer, ""), do: Nx.tensor(0, type: :u32) + defp prompt_token_length(_tokenizer, nil), do: Nx.tensor(0, type: :u32) + + defp prompt_token_length(tokenizer, prompt) do + tokenizer = Bumblebee.configure(tokenizer, length: nil) + inputs = Bumblebee.apply_tokenizer(tokenizer, [prompt]) + length = Nx.axis_size(inputs["input_ids"], 1) + + last_token = + inputs["input_ids"] + |> Nx.slice_along_axis(length - 1, 1, axis: 1) + |> Nx.squeeze() + |> Nx.to_number() + + special_ids = special_token_ids(tokenizer) + + length = + if last_token in special_ids do + length - 1 + else + length + end + + Nx.as_type(length, :u32) + end + + defp special_token_ids(tokenizer) do + tokenizer.special_tokens + |> Map.values() + |> Enum.concat(tokenizer.additional_special_tokens) + |> MapSet.new(&PreTrainedTokenizer.token_to_id(tokenizer, &1)) + end + + defp resolve_prompt(opts, model_info) do + case {opts[:prompt], opts[:prompt_name]} do + {prompt, nil} when is_binary(prompt) -> + prompt + + {prompt, nil} when prompt in [false, nil] -> + if Keyword.has_key?(opts, :prompt) do + "" + else + find_default_prompt(model_info) + end + + {prompt, nil} -> + raise ArgumentError, "expected :prompt to be a string or false, got: #{inspect(prompt)}" + + {nil, prompt_name} when is_binary(prompt_name) -> + find_prompt_by_name!(model_info, prompt_name) + + {nil, prompt_name} when prompt_name in [false, nil] -> + if Keyword.has_key?(opts, :prompt_name) do + "" + else + find_default_prompt(model_info) + end + + {nil, prompt_name} -> + raise ArgumentError, + "expected :prompt_name to be a string or false, got: #{inspect(prompt_name)}" + + {_prompt, _prompt_name} -> + raise ArgumentError, "expected either :prompt or :prompt_name, but both were given" + end + end + + defp find_prompt_by_name!(model_info, prompt_name) do + prompts = get_in(model_info, [:sentence_transformers, :prompts]) || %{} + + case Map.fetch(prompts, prompt_name) do + {:ok, template} -> + template + + :error -> + raise ArgumentError, + "unknown prompt name #{inspect(prompt_name)}. " <> + "Available prompts: #{inspect(Map.keys(prompts))}" + end + end + + defp find_default_prompt(model_info) do + st = Map.get(model_info, :sentence_transformers, %{}) || %{} + + if default_name = st[:default_prompt_name] do + get_in(st, [:prompts, default_name]) + end + end + + @doc """ + Truncates embeddings along the last dimension to `truncate_dim` (Matryoshka representation). + + ## Options + + * `:normalize` - whether to re-apply L2 normalization after truncation. + Defaults to `false`. + + """ + def truncate_embeddings(embeddings, truncate_dim, opts \\ []) + when is_integer(truncate_dim) and truncate_dim > 0 do + normalize? = Keyword.get(opts, :normalize, false) + truncated = Nx.slice_along_axis(embeddings, 0, truncate_dim, axis: -1) + + if normalize? do + Bumblebee.Utils.Nx.normalize(truncated) + else + truncated + end + end + + @doc """ + Loads root prompt metadata and the first module's saved token limit independently. + """ + def load_config(repository, opts \\ []) do + with {:ok, root} <- optional_config(repository, "config_sentence_transformers.json", opts), + {:ok, legacy} <- optional_config(repository, "sentence_bert_config.json", opts), + {:ok, module_path} <- first_module_path(repository, opts), + {:ok, transformer} <- + optional_config(repository, Path.join(module_path, "sentence_bert_config.json"), opts), + {:ok, tokenizer} <- + optional_config(repository, Path.join(module_path, "tokenizer_config.json"), opts), + {:ok, root_tokenizer} <- optional_config(repository, "tokenizer_config.json", opts) do + configs = Enum.reject([transformer, legacy, root], &is_nil/1) + + if configs == [] and tokenizer == nil and root_tokenizer == nil do + {:error, "could not find sentence transformers config"} + else + config = Enum.reduce(configs, %{}, &Map.merge(&2, &1)) + tokenizer_limit = (tokenizer || root_tokenizer || %{})["model_max_length"] + # Hugging Face uses enormous sentinel values for tokenizers without a limit. + tokenizer_limit = + if is_integer(tokenizer_limit) and tokenizer_limit > 0 and + tokenizer_limit < 1_000_000_000, + do: tokenizer_limit + + limit = + (transformer || %{})["max_seq_length"] || config["max_seq_length"] || tokenizer_limit + + {:ok, parse_st_config(Map.put(config, "max_seq_length", limit))} + end + end + end + + defp first_module_path(repository, opts) do + case Downloader.download_file(repository, "modules.json", opts) do + {:ok, path} -> + with {:ok, modules} <- Pipeline.decode_json(path), + [first | _] <- Enum.sort_by(modules, &(&1["idx"] || 0)) do + {:ok, first["path"] || ""} + else + {:error, reason} -> {:error, reason} + _ -> {:error, "modules.json is empty"} + end + + {:error, reason} -> + if Downloader.missing_file?(repository, "modules.json", reason), + do: {:ok, "0_Transformer"}, + else: {:error, reason} + end + end + + defp optional_config(repository, filename, opts) do + case Downloader.download_file(repository, filename, opts) do + {:ok, path} -> + Pipeline.decode_json(path) + + {:error, reason} -> + if Downloader.missing_file?(repository, filename, reason), + do: {:ok, nil}, + else: {:error, reason} + end + end + + defp load_config_and_attach(repository, model_info) do + case load_config(repository) do + {:ok, config} -> + {:ok, Map.put(model_info, :sentence_transformers, config)} + + {:error, "could not find sentence transformers config"} -> + {:ok, Map.put(model_info, :sentence_transformers, default_st_config())} + + {:error, reason} -> + {:error, reason} + end + end + + defp parse_st_config(config) do + %{ + prompts: Map.get(config, "prompts", %{}), + default_prompt_name: config["default_prompt_name"], + similarity_fn_name: parse_similarity_fn_name(config["similarity_fn_name"]), + max_seq_length: config["max_seq_length"] + } + end + + defp default_st_config do + %{ + prompts: %{}, + default_prompt_name: nil, + similarity_fn_name: nil, + max_seq_length: nil + } + end + + defp parse_similarity_fn_name("cosine"), do: :cosine + defp parse_similarity_fn_name("dot"), do: :dot + defp parse_similarity_fn_name("euclidean"), do: :euclidean + defp parse_similarity_fn_name("manhattan"), do: :manhattan + defp parse_similarity_fn_name(nil), do: nil + defp parse_similarity_fn_name(other), do: other + + @doc """ + Infers fallback pooling mode when `modules.json` is missing. + Returns `"last"` for causal-LM architectures, and `"mean"` otherwise. + """ + def infer_fallback_pooling_mode(repository, model_info) do + is_causal_from_config = + case Downloader.download_file(repository, "config.json") do + {:ok, path} -> + case Pipeline.decode_json(path) do + {:ok, config} -> + archs = config["architectures"] || [] + + causal_arch? = + Enum.any?(archs, fn arch -> + is_binary(arch) and + (String.ends_with?(arch, "ForCausalLM") or + String.ends_with?(arch, "LMHeadModel")) + end) + + causal_arch? and Map.get(config, "is_causal", true) + + _ -> + false + end + + _ -> + false + end + + is_causal_from_spec = + if spec = model_info[:spec] || model_info[:config] do + struct_name = spec.__struct__ |> Module.split() |> List.last() + + struct_name in [ + "Gemma", + "GemmaConfig", + "Gemma3Text", + "Gemma3TextConfig", + "Llama", + "LlamaConfig", + "Mistral", + "MistralConfig", + "Qwen3", + "Qwen3Config", + "SmolLm3", + "Smollm3Config", + "Phi", + "PhiConfig", + "Phi3", + "Phi3Config", + "Gpt2", + "Gpt2Config", + "GptBigCode", + "GptBigCodeConfig", + "GptNeoX", + "GptNeoXConfig" + ] + else + false + end + + if is_causal_from_config or is_causal_from_spec do + "last" + else + "mean" + end + end +end diff --git a/lib/bumblebee/sentence_transformers/downloader.ex b/lib/bumblebee/sentence_transformers/downloader.ex new file mode 100644 index 00000000..27241fa5 --- /dev/null +++ b/lib/bumblebee/sentence_transformers/downloader.ex @@ -0,0 +1,108 @@ +defmodule Bumblebee.SentenceTransformers.Downloader do + @moduledoc false + + alias Bumblebee.HuggingFace.Hub + + def subrepository(repository, path) when path in [nil, ""], do: repository + def subrepository({:local, dir}, path), do: {:local, Path.join(dir, path)} + + def subrepository({:local, dir, opts}, path) do + {:local, Path.join([dir, opts[:subdir] || "", path])} + end + + def subrepository({:hf, id}, path), do: {:hf, id, subdir: path} + + def subrepository({:hf, id, opts}, path) do + {:hf, id, Keyword.put(opts, :subdir, Path.join(opts[:subdir] || "", path))} + end + + @doc """ + Downloads or loads a specific file from a repository on demand. + """ + def download_file(repository, filename, opts \\ []) + + def download_file({:local, dir}, filename, _opts) when is_binary(dir) do + path = Path.join(dir, filename) + + if File.exists?(path) do + {:ok, path} + else + {:error, "file #{inspect(path)} does not exist"} + end + end + + def download_file({:local, dir, local_opts}, filename, _opts) when is_binary(dir) do + path = + if subdir = local_opts[:subdir] do + Path.join([dir, subdir, filename]) + else + Path.join(dir, filename) + end + + if File.exists?(path) do + {:ok, path} + else + {:error, "file #{inspect(path)} does not exist"} + end + end + + def download_file({:hf, repo_id}, filename, opts) do + download_file({:hf, repo_id, []}, filename, opts) + end + + def download_file({:hf, repo_id, repo_opts}, filename, opts) do + merged_opts = Keyword.merge(repo_opts, opts) + subdir = merged_opts[:subdir] + + full_filename = + if subdir do + subdir <> "/" <> filename + else + filename + end + + url = Hub.file_url(repo_id, full_filename, merged_opts[:revision]) + cache_scope = repo_id_to_cache_scope(repo_id) + + Hub.cached_download( + url, + [cache_scope: cache_scope] ++ Keyword.take(merged_opts, [:cache_dir, :offline, :auth_token]) + ) + end + + def download_file(other, _filename, _opts) do + {:error, "invalid repository format: #{inspect(other)}"} + end + + @doc false + def missing_file?(repository, filename, reason) do + case repository do + {:local, dir} -> missing_local_file?(Path.join(dir, filename)) + {:local, dir, opts} -> missing_local_file?(Path.join([dir, opts[:subdir] || "", filename])) + _ -> is_binary(reason) and String.starts_with?(reason, "file not found, url: ") + end + end + + defp missing_local_file?(path) do + case File.stat(path) do + {:error, :enoent} -> true + _ -> false + end + end + + @doc """ + Checks whether a file exists in the repository. + """ + def file_exists?(repository, filename, opts \\ []) do + case download_file(repository, filename, opts) do + {:ok, _path} -> true + {:error, _} -> false + end + end + + defp repo_id_to_cache_scope(repo_id) do + repo_id + |> String.replace("/", "--") + |> String.replace(~r/[^\w-]/, "") + end +end diff --git a/lib/bumblebee/sentence_transformers/pipeline.ex b/lib/bumblebee/sentence_transformers/pipeline.ex new file mode 100644 index 00000000..c4b5434b --- /dev/null +++ b/lib/bumblebee/sentence_transformers/pipeline.ex @@ -0,0 +1,816 @@ +defmodule Bumblebee.SentenceTransformers.Pipeline do + @moduledoc false + + alias Axon.MixedPrecision.Policy + alias Bumblebee.Layers + alias Bumblebee.SentenceTransformers.Downloader + + @doc """ + Builds the pipeline modules on top of base model_info. + """ + def build(repository, modules, model_info, opts) do + base_model = model_info.model + + 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) + + with {:ok, modules} <- maybe_fuse_dense(repository, modules, opts[:fuse_dense]) do + has_weighted_pooling? = + Enum.any?(modules, fn m -> + normalize_module_type(m["type"]) == "WeightedLayerPooling" + end) + + extracted_hidden_states = + if has_weighted_pooling?, do: extract_all_hidden_states(base_model) + + hidden_state = + if extracted_hidden_states do + extracted_hidden_states + else + Axon.nx(base_model, fn + %{hidden_state: hidden_state} -> + hidden_state + + output -> + raise ArgumentError, + "expected base model output to contain :hidden_state, but got keys: #{inspect(Map.keys(output))}." <> + " Ensure you loaded the base encoder model." + end) + 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, false} + + result = + Enum.reduce_while( + modules, + {:ok, initial_state}, + fn module, {:ok, {node, params_data, params_parameters, prompt_length_required}} -> + case build_module( + repository, + module, + node, + attention_mask, + params_data, + params_parameters, + opts + ) do + {:ok, node, params_data, params_parameters, requires_prompt_length} -> + {:cont, + {:ok, + {node, params_data, params_parameters, + prompt_length_required or requires_prompt_length}}} + + {:error, reason} -> + {:halt, {:error, reason}} + end + end + ) + + with {:ok, {embedding, params_data, params_parameters, prompt_length_required}} <- result do + embedding = + case opts[:truncate_dim] do + truncate_dim when is_integer(truncate_dim) and truncate_dim > 0 -> + truncated = + Axon.nx( + embedding, + fn tensor -> Nx.slice_along_axis(tensor, 0, truncate_dim, axis: -1) end, + name: "truncate_dim" + ) + + needs_renorm? = + opts[:normalize] == true or + opts[:normalize_embeddings] == true or + pipeline_has_normalize?(modules) + + if needs_renorm? do + Axon.nx(truncated, &Bumblebee.Utils.Nx.normalize/1, + name: "truncate_dim.normalize" + ) + else + truncated + end + + _ -> + embedding + end + + 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, + Map.put_new( + %{model_info | model: final_model, params: final_params}, + :prompt_length_required, + prompt_length_required + )} + end + end + end + + def normalize_module_type(type) when is_binary(type) do + type + |> String.split(".") + |> List.last() + end + + @doc """ + Returns true if the built model expects a `prompt_length` input. + This happens when a pooling module has `include_prompt: false`. + """ + def prompt_length_required?(%{prompt_length_required: value}), do: value + def prompt_length_required?(_), do: false + + defp build_module( + repository, + %{"path" => path} = module, + node, + attention_mask, + params_data, + params_parameters, + opts + ) do + case normalize_module_type(module["type"]) do + "Pooling" -> + config_result = + cond do + config = module["config"] -> + {:ok, config} + + path == "default_pooling" -> + mode = module["pooling_mode"] || "mean" + {:ok, %{"pooling_mode" => mode}} + + true -> + with {:ok, config_path} <- + Downloader.download_file(repository, Path.join(path, "config.json")) do + decode_json(config_path) + end + end + + with {:ok, config} <- config_result, + {:ok, node} <- pooling_layer(node, attention_mask, config, path) do + requires_prompt_length = Map.get(config, "include_prompt", true) == false + {:ok, node, params_data, params_parameters, requires_prompt_length} + end + + "Dense" -> + with {:ok, config_path} <- + Downloader.download_file(repository, Path.join(path, "config.json")), + {:ok, config} <- decode_json(config_path), + {:ok, activation} <- parse_activation(config["activation_function"], path), + {:ok, weights_file, weights_path} <- find_weights_file(repository, path) do + in_features = config["in_features"] + out_features = config["out_features"] + bias = Map.get(config, "bias", true) + use_residual = Map.get(config, "use_residual", false) + layer_name = path + + out_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 + + {node, params_data, params_parameters} = + if use_residual do + if in_features == out_features do + {Axon.add(out_node, node, name: "#{layer_name}.add_residual"), params_data, + params_parameters} + else + res_name = "#{layer_name}.residual" + res_node = Axon.dense(node, out_features, use_bias: false, name: res_name) + + res_kernel = + tensors["residual.weight"] + |> Nx.to_tensor() + |> Nx.transpose() + |> cast_param(opts[:type]) + |> allocate_param(opts[:backend]) + + res_params = %{"kernel" => res_kernel} + params_data = Map.put(params_data, res_name, res_params) + + params_parameters = + if params_parameters do + Map.put(params_parameters, res_name, ["kernel"]) + end + + {Axon.add(out_node, res_node, name: "#{layer_name}.add_residual"), params_data, + params_parameters} + end + else + {out_node, params_data, params_parameters} + end + + {:ok, node, params_data, params_parameters, false} + end + + "FusedDense" -> + [p1, p2] = module["paths"] + layer_name = module["path"] + out_features = module["out_features"] + + with {:ok, w1_file, w1_path} <- find_weights_file(repository, p1), + {:ok, w2_file, w2_path} <- find_weights_file(repository, p2) 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 = + k1 + |> Nx.dot(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, false} + end + + "Normalize" -> + node = Axon.nx(node, &Bumblebee.Utils.Nx.normalize/1, name: "#{path}.normalize") + {:ok, node, params_data, params_parameters, false} + + "Dropout" -> + {:ok, node, params_data, params_parameters, false} + + "LayerNorm" -> + with {:ok, config_path} <- + Downloader.download_file(repository, Path.join(path, "config.json")), + {:ok, config} <- decode_json(config_path), + {:ok, weights_file, weights_path} <- find_weights_file(repository, path) do + eps = config["eps"] || 1.0e-5 + layer_name = path + + node = Axon.layer_norm(node, epsilon: eps, name: layer_name) + tensors = load_tensors(weights_file, weights_path, opts) + + gamma = + (tensors["weight"] || tensors["gamma"] || tensors["norm.weight"]) + |> Nx.to_tensor() + |> cast_param(opts[:type]) + |> allocate_param(opts[:backend]) + + beta = + (tensors["bias"] || tensors["beta"] || tensors["norm.bias"]) + |> Nx.to_tensor() + |> cast_param(opts[:type]) + |> allocate_param(opts[:backend]) + + layer_params = %{"gamma" => gamma, "beta" => beta} + params_data = Map.put(params_data, layer_name, layer_params) + + params_parameters = + if params_parameters do + Map.put(params_parameters, layer_name, ["gamma", "beta"]) + end + + {:ok, node, params_data, params_parameters, false} + end + + "WeightedLayerPooling" -> + with {:ok, config_path} <- + Downloader.download_file(repository, Path.join(path, "config.json")), + {:ok, config} <- decode_json(config_path) do + layer_start = config["layer_start"] || 4 + num_hidden_layers = config["num_hidden_layers"] || 12 + num_weights = max(num_hidden_layers + 1 - layer_start, 1) + layer_name = path + + {_weights, params_data, params_parameters} = + case find_weights_file(repository, path) do + {:ok, weights_file, weights_path} -> + tensors = load_tensors(weights_file, weights_path, opts) + + w = + (tensors["layer_weights"] || tensors["weight"] || + Nx.broadcast(1.0, {num_weights})) + |> Nx.to_tensor() + |> cast_param(opts[:type]) + |> allocate_param(opts[:backend]) + + params_data = Map.put(params_data, layer_name, %{"layer_weights" => w}) + + params_parameters = + if params_parameters do + Map.put(params_parameters, layer_name, ["layer_weights"]) + end + + {w, params_data, params_parameters} + + _ -> + w = + 1.0 + |> Nx.broadcast({num_weights}) + |> cast_param(opts[:type]) + |> allocate_param(opts[:backend]) + + params_data = Map.put(params_data, layer_name, %{"layer_weights" => w}) + + params_parameters = + if params_parameters do + Map.put(params_parameters, layer_name, ["layer_weights"]) + end + + {w, params_data, params_parameters} + end + + param_layer = Axon.param("layer_weights", {num_weights}, initializer: :ones) + + weighted_node = + Axon.layer( + fn stacked_layers, layer_weights, _opts -> + total_layers = Nx.axis_size(stacked_layers, 0) + + selected = + Nx.slice_along_axis( + stacked_layers, + layer_start, + total_layers - layer_start, + axis: 0 + ) + + num_selected = total_layers - layer_start + weights_reshaped = Nx.reshape(layer_weights, {num_selected, 1, 1, 1}) + weighted = Nx.multiply(weights_reshaped, selected) + sum = Nx.sum(weighted, axes: [0]) + total_weight = Nx.sum(layer_weights) + Nx.divide(sum, total_weight) + end, + [node, param_layer], + name: layer_name + ) + + {:ok, weighted_node, params_data, params_parameters, false} + end + + other -> + {:error, "unsupported SentenceTransformers module #{inspect(other)}"} + end + end + + def default_pooling(node, attention_mask, name \\ "default_pooling") do + config = %{"pooling_mode_mean_tokens" => true} + pooling_layer(node, attention_mask, config, name) + end + + def pooling_layer(hidden_state, attention_mask, config, name) do + include_prompt = Map.get(config, "include_prompt", true) + + modes_from_flags = + 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_weightedmean_tokens"], :weightedmean_tokens}, + {["pooling_mode_lasttoken", "pooling_mode_last_token"], :last_token} + ], + Enum.any?(keys, &(config[&1] == true)), + do: mode + + modes = + case {modes_from_flags, config["pooling_mode"]} do + {[], "mean"} -> [:mean_tokens] + {[], "cls"} -> [:cls_token] + {[], "max"} -> [:max_tokens] + {[], "mean_sqrt_len_tokens"} -> [:mean_sqrt_len_tokens] + {[], "weightedmean"} -> [:weightedmean_tokens] + {[], "last"} -> [:last_token] + {[], "lasttoken"} -> [:last_token] + {modes, _} -> modes + end + + if modes == [] do + {:error, "no supported pooling mode found in #{name}"} + else + layer = build_pooling_layer(hidden_state, attention_mask, modes, include_prompt, name) + {:ok, layer} + end + end + + defp build_pooling_layer(hidden_state, attention_mask, modes, true, name) do + Axon.layer( + fn hidden_state, attention_mask, _opts -> + apply_pooling_modes(hidden_state, attention_mask, modes) + end, + [hidden_state, attention_mask], + name: name + ) + end + + defp build_pooling_layer(hidden_state, attention_mask, modes, false, name) do + prompt_length = Axon.input("prompt_length", optional: true) + + Axon.layer( + fn hidden_state, attention_mask, prompt_length, _opts -> + attention_mask = + case prompt_length do + %Axon.None{} -> attention_mask + _ -> exclude_prompt_from_mask(attention_mask, prompt_length) + end + + apply_pooling_modes(hidden_state, attention_mask, modes) + end, + [hidden_state, attention_mask, prompt_length], + name: name + ) + end + + defp apply_pooling_modes(hidden_state, attention_mask, modes) do + 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 = mask |> Nx.sum(axes: [1]) |> Nx.max(eps) + Nx.divide(sum_embeddings, sum_mask) + + :cls_token -> + first_indices = Nx.argmax(attention_mask, axis: 1) + Bumblebee.Utils.Nx.batched_take(hidden_state, first_indices) + + :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.neg_infinity(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 = mask |> Nx.sum(axes: [1]) |> Nx.max(eps) + Nx.divide(sum_embeddings, Nx.sqrt(sum_mask)) + + :weightedmean_tokens -> + mask = attention_mask |> Nx.as_type(type) |> Nx.new_axis(-1) + seq_len = Nx.axis_size(hidden_state, 1) + + weights = + {1, seq_len, 1} + |> Nx.iota(type: type) + |> Nx.add(1) + + weighted_mask = Nx.multiply(mask, weights) + sum_embeddings = Nx.sum(Nx.multiply(hidden_state, weighted_mask), axes: [1]) + sum_mask = weighted_mask |> Nx.sum(axes: [1]) |> Nx.max(eps) + Nx.divide(sum_embeddings, sum_mask) + + :last_token -> + lengths = + attention_mask + |> Nx.not_equal(0) + |> Nx.select( + Nx.iota(Nx.shape(attention_mask), axis: 1), + 0 + ) + |> Nx.reduce_max(axes: [1]) + |> Nx.as_type({:s, 64}) + + selected = Bumblebee.Utils.Nx.batched_take(hidden_state, lengths) + valid = attention_mask |> Nx.not_equal(0) |> Nx.any(axes: [1]) |> Nx.new_axis(-1) + Nx.select(Nx.broadcast(valid, Nx.shape(selected)), selected, 0) + end) + + case pool_outputs do + [single] -> single + multiple -> Nx.concatenate(multiple, axis: -1) + end + end + + defp exclude_prompt_from_mask(attention_mask, prompt_length) do + type = Nx.type(attention_mask) + seq_len = Nx.axis_size(attention_mask, 1) + batch_size = Nx.axis_size(attention_mask, 0) + + prompt_length = + case Nx.rank(prompt_length) do + 0 -> Nx.broadcast(Nx.as_type(prompt_length, type), {batch_size, 1}) + 1 -> prompt_length |> Nx.as_type(type) |> Nx.reshape({batch_size, 1}) + 2 -> Nx.as_type(prompt_length, type) + end + + positions = Nx.iota({1, seq_len}, type: type) + + pad_lengths = + attention_mask + |> Nx.as_type({:s, 64}) + |> Nx.argmax(axis: 1) + |> Nx.as_type(type) + |> Nx.reshape({batch_size, 1}) + + has_left_pad = Nx.any(Nx.not_equal(pad_lengths, 0)) + + threshold_right = prompt_length + threshold_left = Nx.add(pad_lengths, prompt_length) + + threshold = + Nx.select(Nx.broadcast(has_left_pad, {batch_size, 1}), threshold_left, threshold_right) + + keep_mask = Nx.greater_equal(positions, threshold) + + Nx.multiply(attention_mask, Nx.as_type(keep_mask, type)) + 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("torch.nn.modules.activation.Sigmoid", _path), do: {:ok, :sigmoid} + defp parse_activation("torch.nn.modules.activation.Mish", _path), do: {:ok, :mish} + defp parse_activation("torch.nn.modules.activation.LeakyReLU", _path), do: {:ok, :leaky_relu} + defp parse_activation("transformers.activations.GELUActivation", _path), do: {:ok, :gelu} + defp parse_activation("transformers.activations.FastGELUActivation", _path), do: {:ok, :gelu} + defp parse_activation("transformers.activations.NewGELUActivation", _path), do: {:ok, :gelu} + defp parse_activation("transformers.activations.SiLUActivation", _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 pipeline_has_normalize?(modules) do + Enum.any?(modules, fn m -> + normalize_module_type(m["type"]) == "Normalize" + end) + end + + defp maybe_activation(node, nil), do: node + defp maybe_activation(node, activation), do: Axon.activation(node, activation) + + defp maybe_fuse_dense(repository, modules, true) do + fuse_dense_modules(repository, modules, []) + end + + defp maybe_fuse_dense(_repository, modules, _false), do: {:ok, modules} + + defp fuse_dense_modules( + repository, + [%{"type" => t1, "path" => p1} = m1, %{"type" => t2, "path" => p2} = m2 | rest], + acc + ) do + if normalize_module_type(t1) == "Dense" and normalize_module_type(t2) == "Dense" do + with {:ok, c1} <- load_module_config(repository, p1), + {:ok, c2} <- load_module_config(repository, p2), + 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(repository, rest, [fused_module | acc]) + else + _ -> + fuse_dense_modules(repository, [m2 | rest], [m1 | acc]) + end + else + fuse_dense_modules(repository, [m2 | rest], [m1 | acc]) + end + end + + defp fuse_dense_modules(repository, [other | rest], acc) do + fuse_dense_modules(repository, rest, [other | acc]) + end + + defp fuse_dense_modules(_repository, [], acc) do + {:ok, Enum.reverse(acc)} + end + + defp load_module_config(repository, path) do + with {:ok, config_path} <- + Downloader.download_file(repository, Path.join(path, "config.json")) do + decode_json(config_path) + 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 + + @weights_filenames ["model.safetensors", "pytorch_model.bin"] + + defp find_weights_file(repository, dir) do + Enum.find_value(@weights_filenames, fn filename -> + path = Path.join(dir, filename) + + case Downloader.download_file(repository, path) do + {:ok, downloaded_path} -> {:ok, path, downloaded_path} + _ -> nil + end + end) || {:error, "could not find parameters file in #{dir}"} + 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 + + @doc false + def cast_param(tensor, nil), do: tensor + + def cast_param(tensor, %Policy{params: type}) do + Nx.as_type(tensor, type) + end + + def cast_param(tensor, type) do + type = Nx.Type.normalize!(type) + Nx.as_type(tensor, type) + end + + @doc false + def allocate_param(tensor, nil), do: tensor + + def 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, %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 + + def 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 + + @doc """ + Builds a base model from a StaticEmbedding module (e.g. Model2Vec). + """ + def build_static_embedding(repository, %{"path" => path}, opts) do + with {:ok, weights_file, weights_path} <- find_weights_file(repository, path) do + tensors = load_tensors(weights_file, weights_path, opts) + + embedding_weight = + (tensors["embedding.weight"] || tensors["weight"]) + |> Nx.to_tensor() + |> cast_param(opts[:type]) + |> allocate_param(opts[:backend]) + + {vocab_size, embedding_dim} = Nx.shape(embedding_weight) + layer_name = "#{path}.embedding" + + model = + "input_ids" + |> Axon.input(optional: false) + |> Axon.embedding(vocab_size, embedding_dim, name: layer_name) + |> Axon.nx(fn hidden_state -> + %{hidden_state: hidden_state} + end) + + params = %Axon.ModelState{ + data: %{layer_name => %{"kernel" => embedding_weight}}, + parameters: %{layer_name => ["kernel"]}, + frozen_parameters: %{}, + state: %{} + } + + {:ok, %{model: model, params: params, spec: nil}} + end + end + + defp extract_all_hidden_states(%Axon{nodes: nodes}) do + opt_in_node = + Enum.find_value(nodes, fn {_id, node} -> + if node.op_name == :global_opt_in and :output_hidden_states in node.global_options do + node + end + end) + + if opt_in_node do + [raw_id] = opt_in_node.parent + raw_axon = %Axon{output: raw_id, nodes: nodes} + + Axon.nx( + raw_axon, + fn tuple_hs -> + tuple_hs |> Tuple.to_list() |> Nx.stack(axis: 0) + end, + name: "stacked_hidden_states" + ) + end + end + + defp extract_all_hidden_states(_), do: nil +end diff --git a/mix.exs b/mix.exs index 5d84035b..aa410b91 100644 --- a/mix.exs +++ b/mix.exs @@ -73,6 +73,10 @@ defmodule Bumblebee.MixProject do Bumblebee.Diffusion.StableDiffusion, Bumblebee.Diffusion.StableDiffusionControlNet ], + "Sentence Transformers": [ + Bumblebee.SentenceTransformers, + Bumblebee.SentenceTransformers.Pipeline + ], Models: [ Bumblebee.Audio.Whisper, Bumblebee.Diffusion.ControlNet, diff --git a/test/bumblebee/sentence_transformers/dense_edge_test.exs b/test/bumblebee/sentence_transformers/dense_edge_test.exs new file mode 100644 index 00000000..f519cca2 --- /dev/null +++ b/test/bumblebee/sentence_transformers/dense_edge_test.exs @@ -0,0 +1,260 @@ +defmodule Bumblebee.SentenceTransformers.DenseEdgeTest do + use ExUnit.Case, async: true + + import Bumblebee.TestHelpers + + alias Bumblebee.SentenceTransformers + alias Bumblebee.SentenceTransformers.Downloader + alias Bumblebee.SentenceTransformers.Pipeline + + test "last-token pooling handles left padding, prompt exclusion, and empty masks" do + hidden = Axon.input("hidden", shape: {nil, nil, 2}) + mask = Axon.input("attention_mask", shape: {nil, nil}) + + assert {:ok, model} = + Pipeline.pooling_layer( + hidden, + mask, + %{"pooling_mode_lasttoken" => true, "include_prompt" => false}, + "last" + ) + + inputs = %{ + "hidden" => + Nx.tensor([ + [[99.0, 99.0], [1.0, 2.0], [3.0, 4.0]], + [[5.0, 6.0], [7.0, 8.0], [99.0, 99.0]] + ]), + "attention_mask" => Nx.tensor([[0, 1, 1], [1, 1, 0]]), + "prompt_length" => Nx.tensor([1, 1]) + } + + {init, predict} = Axon.build(model) + params = init.(inputs, Axon.ModelState.empty()) + assert_all_close(predict.(params, inputs), Nx.tensor([[3.0, 4.0], [7.0, 8.0]])) + end + + @tag :tmp_dir + test "root prompts and arbitrary first-module token limits are merged", %{tmp_dir: dir} do + File.mkdir_p!(Path.join(dir, "encoder")) + + File.write!( + Path.join(dir, "modules.json"), + Jason.encode!([%{"idx" => 0, "path" => "encoder", "type" => "Transformer"}]) + ) + + File.write!( + Path.join(dir, "config_sentence_transformers.json"), + Jason.encode!(%{"prompts" => %{"task" => "hello "}, "default_prompt_name" => "task"}) + ) + + File.write!( + Path.join(dir, "encoder/sentence_bert_config.json"), + Jason.encode!(%{"max_seq_length" => 3}) + ) + + assert {:ok, config} = SentenceTransformers.load_config({:local, dir}) + assert config.max_seq_length == 3 + assert config.default_prompt_name == "task" + assert config.prompts == %{"task" => "hello "} + + File.rm!(Path.join(dir, "encoder/sentence_bert_config.json")) + + File.write!( + Path.join(dir, "encoder/tokenizer_config.json"), + Jason.encode!(%{"model_max_length" => 5}) + ) + + assert {:ok, config} = SentenceTransformers.load_config({:local, dir}) + assert config.max_seq_length == 5 + end + + @tag :tmp_dir + test "saved truncation stays separate from compilation padding and explicit overrides", %{ + tmp_dir: dir + } do + tokenizer = tiny_tokenizer(dir) + ids = Axon.input("input_ids", shape: {nil, nil}) + mask = Axon.input("attention_mask", shape: {nil, nil}) + + model = + Axon.layer( + fn ids, mask, _opts -> + ids |> Nx.multiply(mask) |> Nx.sum(axes: [1], keep_axes: true) |> Nx.as_type(:f32) + end, + [ids, mask] + ) + + info = %{ + model: model, + params: Axon.ModelState.empty(), + spec: nil, + sentence_transformers: %{max_seq_length: 3} + } + + text = "hello world hello world hello world" + + for opts <- [ + [], + [compile: [batch_size: 1, sequence_length: 8]], + [compile: [batch_size: 1, sequence_length: [4, 8]]] + ] do + assert %{embedding: embedding} = + Nx.Serving.run(SentenceTransformers.text_embedding(info, tokenizer, opts), text) + + assert Nx.to_flat_list(embedding) == [4.0] + + assert %{embedding: embedding} = + Nx.Serving.run( + SentenceTransformers.text_embedding( + info, + tokenizer, + Keyword.put(opts, :prompt, "world ") + ), + text + ) + + assert Nx.to_flat_list(embedding) == [5.0] + end + + assert %{embedding: embedding} = + Nx.Serving.run( + SentenceTransformers.text_embedding(info, tokenizer, max_seq_length: 5), + text + ) + + assert Nx.to_flat_list(embedding) == [7.0] + + configured = + SentenceTransformers.configure_embedding_tokenizer(info, tokenizer, + compile: [batch_size: 1, sequence_length: 8] + ) + + tokens = SentenceTransformers.apply_embedding_tokenizer(configured, [text], 8) + assert Nx.to_list(tokens["attention_mask"]) == [[1, 1, 1, 0, 0, 0, 0, 0]] + end + + @tag :tmp_dir + test "saved limits truncate without forcing padding or bypassing smaller buckets", %{ + tmp_dir: dir + } do + tokenizer = tiny_tokenizer(dir) + info = %{sentence_transformers: %{max_seq_length: 512}} + + dynamic = SentenceTransformers.configure_embedding_tokenizer(info, tokenizer, []) + + assert Nx.shape( + SentenceTransformers.apply_embedding_tokenizer(dynamic, ["hello"], nil)["input_ids"] + ) == {1, 1} + + compiled = + SentenceTransformers.configure_embedding_tokenizer(info, tokenizer, + compile: [batch_size: 1, sequence_length: [8, 512]] + ) + + inputs = SentenceTransformers.apply_embedding_tokenizer(compiled, ["hello"], [8, 512]) + assert Nx.shape(inputs["input_ids"]) == {1, 8} + assert Nx.to_list(inputs["attention_mask"]) == [[1, 0, 0, 0, 0, 0, 0, 0]] + + long_text = Enum.join(List.duplicate("hello", 600), " ") + + assert Nx.shape( + SentenceTransformers.apply_embedding_tokenizer(dynamic, [long_text], nil)[ + "input_ids" + ] + ) == + {1, 512} + end + + @tag :tmp_dir + test "compiled padding supports EOS-only tokenizers", %{tmp_dir: dir} do + tokenizer = %{tiny_tokenizer(dir) | special_tokens: %{eos: "[PAD]", unk: "[UNK]"}} + info = %{sentence_transformers: %{max_seq_length: 512}} + + for direction <- [:left, :right] do + tokenizer = %{tokenizer | pad_direction: direction} + + configured = + SentenceTransformers.configure_embedding_tokenizer(info, tokenizer, + compile: [batch_size: 1, sequence_length: [8, 512]] + ) + + inputs = SentenceTransformers.apply_embedding_tokenizer(configured, ["hello"], [8, 512]) + + expected = + if direction == :left, do: [0, 0, 0, 0, 0, 0, 0, 1], else: [1, 0, 0, 0, 0, 0, 0, 0] + + assert Nx.to_list(inputs["input_ids"]) == [expected] + assert Nx.to_list(inputs["attention_mask"]) == [expected] + end + end + + @tag :tmp_dir + test "malformed manifests and download failures propagate from model and tokenizer loading", %{ + tmp_dir: dir + } do + File.write!(Path.join(dir, "modules.json"), "invalid json") + assert {:error, _} = SentenceTransformers.load_model({:local, dir}) + assert {:error, _} = SentenceTransformers.load_tokenizer({:local, dir}) + + assert {:error, "invalid repository format: :invalid"} = + SentenceTransformers.load_model(:invalid) + + assert {:error, "invalid repository format: :invalid"} = + SentenceTransformers.load_tokenizer(:invalid) + + refute Downloader.missing_file?( + {:hf, "repo"}, + "modules.json", + "HTTP request failed with status 500, url: test" + ) + + refute Downloader.missing_file?( + {:hf, "repo"}, + "modules.json", + "repository not found, url: test" + ) + + refute Downloader.missing_file?({:hf, "repo"}, "modules.json", "cache file not found") + + assert Downloader.missing_file?( + {:hf, "repo"}, + "modules.json", + "file not found, url: test" + ) + + assert Downloader.missing_file?({:local, dir}, "absent.json", "missing") + end + + defp tiny_tokenizer(dir) do + path = Path.join(dir, "tokenizer.json") + + File.write!( + path, + Jason.encode!(%{ + "version" => "1.0", + "truncation" => nil, + "padding" => nil, + "added_tokens" => [], + "normalizer" => nil, + "pre_tokenizer" => %{"type" => "Whitespace"}, + "post_processor" => nil, + "decoder" => nil, + "model" => %{ + "type" => "WordLevel", + "vocab" => %{"[PAD]" => 0, "hello" => 1, "world" => 2, "[UNK]" => 3}, + "unk_token" => "[UNK]" + } + }) + ) + + {:ok, native} = Tokenizers.Tokenizer.from_file(path) + + %Bumblebee.Text.PreTrainedTokenizer{ + native_tokenizer: native, + type: :bert, + special_tokens: %{pad: "[PAD]", unk: "[UNK]"}, + add_special_tokens: false + } + end +end diff --git a/test/bumblebee/sentence_transformers/python_comparison_test.exs b/test/bumblebee/sentence_transformers/python_comparison_test.exs new file mode 100644 index 00000000..f98e103b --- /dev/null +++ b/test/bumblebee/sentence_transformers/python_comparison_test.exs @@ -0,0 +1,557 @@ +defmodule Bumblebee.SentenceTransformers.PythonComparisonTest do + use ExUnit.Case, async: false + + import Bumblebee.TestHelpers + + alias Bumblebee.SentenceTransformers + alias Bumblebee.SentenceTransformers.Pipeline + alias Bumblebee.SentenceTransformers.TestSupport.PythonBridge + + @moduletag :slow + @moduletag :python + + test "saved short token limit matches Python with root config, prompts and compilation" do + text = Enum.join(List.duplicate("hello world", 12), " ") + + for prompt <- [nil, "Represent the query: "] do + assert {:ok, %{dir: dir, python_output: expected}} = + PythonBridge.generate_random_sentence_transformer( + max_seq_length: 8, + text: text, + prompt: prompt + ) + + try do + assert {:ok, info} = SentenceTransformers.load_model({:local, dir}) + assert info.sentence_transformers.max_seq_length == 8 + assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) + + for opts <- [[], [compile: [batch_size: 1, sequence_length: 32]]] do + actual = + Nx.Serving.run(SentenceTransformers.text_embedding(info, tokenizer, opts), text).embedding + + assert_all_close(actual, Nx.squeeze(expected), atol: 1.0e-4) + end + after + File.rm_rf!(dir) + end + end + end + + test "pooling matches Python with left padding and fully masked sequences" do + hidden = {3, 5, 4} |> Nx.iota(type: :f32) |> Nx.add(1) + mask = Nx.tensor([[0, 0, 1, 1, 1], [1, 1, 0, 0, 0], [0, 0, 0, 0, 0]]) + modes = ["mean", "cls", "max", "mean_sqrt_len_tokens", "weightedmean", "lasttoken"] + + code = """ + import json + import torch + from sentence_transformers.models import Pooling + hidden = torch.tensor(#{Jason.encode!(Nx.to_list(hidden))}) + mask = torch.tensor(#{Jason.encode!(Nx.to_list(mask))}) + outputs = {} + for mode in #{inspect(modes)}: + result = Pooling(4, pooling_mode=mode)({"token_embeddings": hidden.clone(), "attention_mask": mask.clone()})["sentence_embedding"] + outputs[mode] = {"values": torch.nan_to_num(result, neginf=0).tolist(), "negative_infinity": torch.isneginf(result).int().tolist()} + print("__RESULT__:" + json.dumps(outputs)) + """ + + expected = PythonBridge.run_python_code(code, ["sentence-transformers", "torch"], type: :raw) + + for mode <- modes do + assert {:ok, model} = + Pipeline.pooling_layer( + Axon.input("hidden"), + Axon.input("mask"), + %{"pooling_mode" => mode}, + "pool" + ) + + {init, predict} = Axon.build(model) + inputs = %{"hidden" => hidden, "mask" => mask} + params = init.(inputs, Axon.ModelState.empty()) + actual = predict.(params, inputs) + infinite = Nx.is_infinity(actual) + assert Nx.to_list(infinite) == expected[mode]["negative_infinity"] + + assert_all_close(Nx.select(infinite, 0, actual), Nx.tensor(expected[mode]["values"]), + atol: 1.0e-4 + ) + end + end + + test "runtime prompts override saved defaults and work with compiled prompt exclusion" do + assert {:ok, %{dir: dir}} = + PythonBridge.generate_random_sentence_transformer( + prompt: "Represent the query: ", + include_prompt: false + ) + + try do + assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) + assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) + + for {prompt, compile} <- [ + {"Represent the document: ", nil}, + {"Represent the document: ", [batch_size: 2, sequence_length: 32]}, + {"", nil} + ] do + opts = [prompt: prompt] + opts = if compile, do: Keyword.put(opts, :compile, compile), else: opts + + actual = + model_info + |> SentenceTransformers.text_embedding(tokenizer, opts) + |> Nx.Serving.run("Hello world") + + expected = PythonBridge.run_sentence_transformers(dir, "Hello world", prompt: prompt) + assert_all_close(actual.embedding, expected, atol: 1.0e-4) + end + after + File.rm_rf!(dir) + end + end + + test "query and corpus helpers fall back when no named prompts are saved" do + assert {:ok, %{dir: dir, python_output: expected}} = + PythonBridge.generate_random_sentence_transformer() + + try do + assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) + assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) + + for helper <- [:encode_queries, :encode_corpus] do + [actual] = apply(SentenceTransformers, helper, [model_info, tokenizer, ["Hello world"]]) + assert_all_close(actual.embedding, expected, atol: 1.0e-4) + end + after + File.rm_rf!(dir) + end + end + + describe "On-the-fly generated random SentenceTransformers comparison" do + @pooling_modes [ + "mean_tokens", + "cls_token", + "max_tokens", + "mean_sqrt_len_tokens", + "weightedmean_tokens", + "last_token" + ] + + for mode <- @pooling_modes do + @mode mode + + test "matches Python for random model with #{@mode} pooling" do + assert {:ok, %{dir: dir, python_output: python_output, text: text}} = + PythonBridge.generate_random_sentence_transformer( + pooling_mode: @mode, + with_normalize: false + ) + + try do + assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) + assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) + + serving = SentenceTransformers.text_embedding(model_info, tokenizer) + res = Nx.Serving.run(serving, text) + + assert Nx.shape(res.embedding) == Nx.shape(python_output) + assert_all_close(res.embedding, python_output, atol: 1.0e-4) + after + File.rm_rf!(dir) + end + end + end + + test "matches Python for random model with mean pooling and normalization" do + assert {:ok, %{dir: dir, python_output: python_output, text: text}} = + PythonBridge.generate_random_sentence_transformer( + pooling_mode: "mean_tokens", + with_normalize: true + ) + + try do + assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) + assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) + + serving = SentenceTransformers.text_embedding(model_info, tokenizer) + res = Nx.Serving.run(serving, text) + + assert Nx.shape(res.embedding) == Nx.shape(python_output) + assert_all_close(res.embedding, python_output, atol: 1.0e-4) + after + File.rm_rf!(dir) + end + end + + test "matches Python for random model with Dense projection layer" do + assert {:ok, %{dir: dir, python_output: python_output, text: text}} = + PythonBridge.generate_random_sentence_transformer( + pooling_mode: "mean_tokens", + with_dense: true, + with_normalize: true + ) + + try do + assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) + assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) + + serving = SentenceTransformers.text_embedding(model_info, tokenizer) + res = Nx.Serving.run(serving, text) + + assert Nx.shape(res.embedding) == Nx.shape(python_output) + assert_all_close(res.embedding, python_output, atol: 1.0e-4) + after + File.rm_rf!(dir) + end + end + + test "matches Python for random model with prompt and include_prompt=false" do + prompt = "Represent the sentence for retrieval: " + + assert {:ok, %{dir: dir, python_output: python_output, text: text}} = + PythonBridge.generate_random_sentence_transformer( + pooling_mode: "mean_tokens", + prompt: prompt, + include_prompt: false, + text: "Hello world" + ) + + try do + assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) + assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) + + serving = SentenceTransformers.text_embedding(model_info, tokenizer) + res = Nx.Serving.run(serving, text) + + assert Nx.shape(res.embedding) == Nx.shape(python_output) + assert_all_close(res.embedding, python_output, atol: 1.0e-4) + after + File.rm_rf!(dir) + end + end + + test "matches Python for random model with prompt and include_prompt=true" do + prompt = "Represent the sentence for retrieval: " + + assert {:ok, %{dir: dir, python_output: python_output, text: text}} = + PythonBridge.generate_random_sentence_transformer( + pooling_mode: "mean_tokens", + prompt: prompt, + include_prompt: true, + text: "Hello world" + ) + + try do + assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) + assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) + + serving = SentenceTransformers.text_embedding(model_info, tokenizer) + res = Nx.Serving.run(serving, text) + + assert Nx.shape(res.embedding) == Nx.shape(python_output) + assert_all_close(res.embedding, python_output, atol: 1.0e-4) + after + File.rm_rf!(dir) + end + end + + test "matches Python encode_queries using config_sentence_transformers.json prompt" do + prompt = "Represent the query for retrieval: " + text = "Hello world" + + assert {:ok, %{dir: dir, python_output: python_output}} = + PythonBridge.generate_random_sentence_transformer( + pooling_mode: "mean_tokens", + prompt: prompt, + prompt_name: "query", + text: text + ) + + try do + assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) + assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) + + [res] = SentenceTransformers.encode_queries(model_info, tokenizer, [text]) + + assert Nx.shape(res.embedding) == Nx.shape(python_output) + assert_all_close(res.embedding, python_output, atol: 1.0e-4) + after + File.rm_rf!(dir) + end + end + + test "matches Python encode_corpus using config_sentence_transformers.json prompt" do + query_prompt = "Represent the query for retrieval: " + document_prompt = "Represent the document for retrieval: " + text = "Hello world" + + assert {:ok, %{dir: dir, python_output: python_output}} = + PythonBridge.generate_random_sentence_transformer( + pooling_mode: "mean_tokens", + prompt: query_prompt, + document_prompt: document_prompt, + prompt_name: "document", + text: text + ) + + try do + assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) + assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) + + [res] = SentenceTransformers.encode_corpus(model_info, tokenizer, [text]) + + assert Nx.shape(res.embedding) == Nx.shape(python_output) + assert_all_close(res.embedding, python_output, atol: 1.0e-4) + after + File.rm_rf!(dir) + end + end + end + + describe "On-the-fly generated Dense / Normalize / LayerNorm / Dropout comparison" do + test "matches Python for Dense with Tanh activation" do + text = "Hello world" + + assert {:ok, %{dir: dir, python_output: python_output}} = + PythonBridge.generate_random_sentence_transformer( + pooling_mode: "mean_tokens", + with_dense: true, + dense_activation: "Tanh", + text: text + ) + + try do + assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) + assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) + + serving = SentenceTransformers.text_embedding(model_info, tokenizer) + res = Nx.Serving.run(serving, text) + + assert Nx.shape(res.embedding) == Nx.shape(python_output) + assert_all_close(res.embedding, python_output, atol: 1.0e-4) + after + File.rm_rf!(dir) + end + end + + for activation <- ["Identity", "ReLU", "GELU", "SiLU", "Sigmoid", "Mish", "LeakyReLU"] do + @activation activation + + test "matches Python for Dense with #{@activation} activation" do + text = "Hello world" + + assert {:ok, %{dir: dir, python_output: python_output}} = + PythonBridge.generate_random_sentence_transformer( + pooling_mode: "mean_tokens", + with_dense: true, + dense_activation: @activation, + text: text + ) + + try do + assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) + assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) + + serving = SentenceTransformers.text_embedding(model_info, tokenizer) + res = Nx.Serving.run(serving, text) + + assert Nx.shape(res.embedding) == Nx.shape(python_output) + assert_all_close(res.embedding, python_output, atol: 1.0e-4) + after + File.rm_rf!(dir) + end + end + end + + test "matches Python for Normalize module" do + text = "Hello world" + + assert {:ok, %{dir: dir, python_output: python_output}} = + PythonBridge.generate_random_sentence_transformer( + pooling_mode: "mean_tokens", + with_normalize: true, + text: text + ) + + try do + assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) + assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) + + serving = SentenceTransformers.text_embedding(model_info, tokenizer) + res = Nx.Serving.run(serving, text) + + assert Nx.shape(res.embedding) == Nx.shape(python_output) + assert_all_close(res.embedding, python_output, atol: 1.0e-4) + after + File.rm_rf!(dir) + end + end + + test "matches Python for LayerNorm module" do + text = "Hello world" + + assert {:ok, %{dir: dir, python_output: python_output}} = + PythonBridge.generate_random_sentence_transformer( + pooling_mode: "mean_tokens", + with_layernorm: true, + text: text + ) + + try do + assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) + assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) + + serving = SentenceTransformers.text_embedding(model_info, tokenizer) + res = Nx.Serving.run(serving, text) + + assert Nx.shape(res.embedding) == Nx.shape(python_output) + assert_all_close(res.embedding, python_output, atol: 1.0e-4) + after + File.rm_rf!(dir) + end + end + + test "matches Python for Dropout module in eval mode" do + text = "Hello world" + + assert {:ok, %{dir: dir, python_output: python_output}} = + PythonBridge.generate_random_sentence_transformer( + pooling_mode: "mean_tokens", + with_dropout: true, + text: text + ) + + try do + assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) + assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) + + serving = SentenceTransformers.text_embedding(model_info, tokenizer) + res = Nx.Serving.run(serving, text) + + assert Nx.shape(res.embedding) == Nx.shape(python_output) + assert_all_close(res.embedding, python_output, atol: 1.0e-4) + after + File.rm_rf!(dir) + end + end + + test "matches Python when fusing consecutive linear Dense layers" do + text = "Hello world" + + assert {:ok, %{dir: dir, python_output: python_output}} = + PythonBridge.generate_random_sentence_transformer( + pooling_mode: "mean_tokens", + with_fused_dense: true, + text: text + ) + + try do + assert {:ok, model_info} = + SentenceTransformers.load_model({:local, dir}, fuse_dense: true) + + assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) + + serving = SentenceTransformers.text_embedding(model_info, tokenizer) + res = Nx.Serving.run(serving, text) + + assert Nx.shape(res.embedding) == Nx.shape(python_output) + assert_all_close(res.embedding, python_output, atol: 1.0e-4) + after + File.rm_rf!(dir) + end + end + + for {description, opts} <- [ + {"first layer has bias", [fused_dense_first_bias: true]}, + {"first layer has activation", [fused_dense_first_activation: "Tanh"]}, + {"second layer has bias", [fused_dense_second_bias: true]}, + {"second layer has activation", [fused_dense_second_activation: "Tanh"]} + ] do + @description description + @opts opts + + test "refuses to fuse when #{@description}" do + text = "Hello world" + + assert {:ok, %{dir: dir, python_output: python_output}} = + PythonBridge.generate_random_sentence_transformer( + [pooling_mode: "mean_tokens", with_fused_dense: true, text: text] ++ @opts + ) + + try do + assert {:ok, model_info} = + SentenceTransformers.load_model({:local, dir}, fuse_dense: true) + + assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) + + serving = SentenceTransformers.text_embedding(model_info, tokenizer) + res = Nx.Serving.run(serving, text) + + assert Nx.shape(res.embedding) == Nx.shape(python_output) + assert_all_close(res.embedding, python_output, atol: 1.0e-4) + after + File.rm_rf!(dir) + end + end + end + end + + describe "HuggingFace Hub SentenceTransformers comparison" do + @tag :network + test "matches Python for sentence-transformers/all-MiniLM-L6-v2" do + repo_id = "sentence-transformers/all-MiniLM-L6-v2" + input_text = "Hello world" + + # 1. Elixir SentenceTransformers execution + assert {:ok, model_info} = SentenceTransformers.load_model({:hf, repo_id}) + assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:hf, repo_id}) + + serving = SentenceTransformers.text_embedding(model_info, tokenizer) + elixir_res = Nx.Serving.run(serving, input_text) + + # 2. Python sentence_transformers execution + python_res = PythonBridge.run_sentence_transformers(repo_id, input_text) + + # 3. Assertions and comparison + assert Nx.shape(elixir_res.embedding) == Nx.shape(python_res) + assert_all_close(elixir_res.embedding, python_res, atol: 1.0e-4) + end + end + + describe "HuggingFace tiny-random base models comparison" do + @tiny_random_models [ + {"hf-internal-testing/tiny-random-BertModel", "last_hidden_state"}, + {"hf-internal-testing/tiny-random-RobertaModel", "last_hidden_state"} + ] + + for {repo_id, output_key} <- @tiny_random_models do + @repo_id repo_id + @output_key output_key + + @tag :network + test "matches Python transformers for #{@repo_id}" do + inputs = %{ + "input_ids" => [[10, 20, 30, 40, 50]], + "attention_mask" => [[1, 1, 1, 1, 1]] + } + + elixir_inputs = %{ + "input_ids" => Nx.tensor(inputs["input_ids"]), + "attention_mask" => Nx.tensor(inputs["attention_mask"]) + } + + assert {:ok, %{model: model, params: params}} = Bumblebee.load_model({:hf, @repo_id}) + elixir_outputs = Axon.predict(model, params, elixir_inputs) + + python_tensor = + PythonBridge.run_hf_model(@repo_id, inputs, output_key: @output_key) + + elixir_tensor = Map.fetch!(elixir_outputs, :hidden_state) + assert Nx.shape(elixir_tensor) == Nx.shape(python_tensor) + assert_all_close(elixir_tensor, python_tensor, atol: 1.0e-4) + end + end + end +end diff --git a/test/bumblebee/sentence_transformers/python_environment_test.exs b/test/bumblebee/sentence_transformers/python_environment_test.exs new file mode 100644 index 00000000..22262437 --- /dev/null +++ b/test/bumblebee/sentence_transformers/python_environment_test.exs @@ -0,0 +1,48 @@ +defmodule Bumblebee.SentenceTransformers.PythonEnvironmentTest do + use ExUnit.Case, async: false + + alias Bumblebee.SentenceTransformers.TestSupport.PythonBridge + + @tag :python + @tag :tmp_dir + test "incompatible Python is checked once and reports an actionable version error", %{ + tmp_dir: dir + } do + python = System.find_executable("python3") + assert python, "python3 is needed for the Python environment validation regression" + wrapper = Path.join(dir, "incompatible-python") + count_file = Path.join(dir, "invocations") + + File.write!(wrapper, """ + #!#{python} + import importlib.metadata, sys + with open(#{Jason.encode!(count_file)}, "a") as counter: + counter.write("checked\\n") + importlib.metadata.version = lambda package: "0.0.0" + exec(sys.argv[-1]) + """) + + File.chmod!(wrapper, 0o755) + previous = System.get_env("SENTENCE_TRANSFORMERS_PYTHON") || System.get_env("STING_PYTHON") + System.put_env("SENTENCE_TRANSFORMERS_PYTHON", wrapper) + + on_exit(fn -> + if previous, + do: System.put_env("SENTENCE_TRANSFORMERS_PYTHON", previous), + else: System.delete_env("SENTENCE_TRANSFORMERS_PYTHON") + end) + + for _ <- 1..2 do + assert_raise RuntimeError, + ~r/expected sentence-transformers==6.1.0, found 0.0.0.*(SENTENCE_TRANSFORMERS_PYTHON|STING_PYTHON)/, + fn -> + PythonBridge.run_python_code( + "raise RuntimeError('must not execute')", + ["sentence-transformers"] + ) + end + end + + assert File.read!(count_file) == "checked\n" + end +end diff --git a/test/bumblebee/sentence_transformers_test.exs b/test/bumblebee/sentence_transformers_test.exs new file mode 100644 index 00000000..4c2a209f --- /dev/null +++ b/test/bumblebee/sentence_transformers_test.exs @@ -0,0 +1,1289 @@ +defmodule Bumblebee.SentenceTransformersTest do + use ExUnit.Case, async: true + + alias Bumblebee.SentenceTransformers + + defp assert_all_close(a, b, opts \\ []) do + atol = opts[:atol] || 1.0e-4 + diff = Nx.abs(Nx.subtract(a, b)) + max_diff = diff |> Nx.reduce_max() |> Nx.to_number() + assert max_diff <= atol, "expected tensors to be close, got max diff #{max_diff} > #{atol}" + end + + setup do + model = + "input_ids" + |> Axon.input(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 = + 1.0 + |> Nx.broadcast({batch_size, seq_len, 4}) + |> Nx.as_type({:f, 32}) + + %{hidden_state: hidden_state} + end) + + params = %Axon.ModelState{data: %{}, parameters: %{}} + model_info = %{model: model, params: params, spec: nil} + + [model_info: model_info] + end + + describe "pipeline modules" do + @tag :tmp_dir + test "loads modules.json and builds pooling + dense + normalize pipeline", %{ + model_info: model_info, + tmp_dir: tmp_dir + } do + modules = [ + %{ + "idx" => 0, + "name" => "0", + "path" => "", + "type" => "sentence_transformers.models.Transformer" + }, + %{ + "idx" => 1, + "name" => "1", + "path" => "1_Pooling", + "type" => "sentence_transformers.models.Pooling" + }, + %{ + "idx" => 2, + "name" => "2", + "path" => "2_Dense", + "type" => "sentence_transformers.models.Dense" + }, + %{ + "idx" => 3, + "name" => "3", + "path" => "3_Normalize", + "type" => "sentence_transformers.models.Normalize" + } + ] + + File.write!(Path.join(tmp_dir, "modules.json"), Jason.encode!(modules)) + + # 1_Pooling config + File.mkdir_p!(Path.join(tmp_dir, "1_Pooling")) + pooling_config = %{"pooling_mode_mean_tokens" => true} + File.write!(Path.join(tmp_dir, "1_Pooling/config.json"), Jason.encode!(pooling_config)) + + # 2_Dense config and weights + File.mkdir_p!(Path.join(tmp_dir, "2_Dense")) + + dense_config = %{ + "in_features" => 4, + "out_features" => 2, + "bias" => true, + "activation_function" => "torch.nn.modules.activation.Tanh" + } + + File.write!(Path.join(tmp_dir, "2_Dense/config.json"), Jason.encode!(dense_config)) + + weight = Nx.tensor([[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0]]) + bias = Nx.tensor([0.1, -0.1]) + tensors = %{"linear.weight" => weight, "linear.bias" => bias} + Safetensors.write!(Path.join(tmp_dir, "2_Dense/model.safetensors"), tensors) + + # 3_Normalize (no extra config file needed) + File.mkdir_p!(Path.join(tmp_dir, "3_Normalize")) + + assert {:ok, new_model_info} = + SentenceTransformers.load_embedding_head({:local, tmp_dir}, model_info) + + {_init, predict} = Axon.build(new_model_info.model) + + inputs = %{ + "input_ids" => Nx.tensor([[1, 2]]), + "attention_mask" => Nx.tensor([[1, 1]]) + } + + output = predict.(new_model_info.params, inputs) + + assert %{embedding: embedding, pooled_state: pooled, hidden_state: _} = output + assert Nx.shape(embedding) == {1, 2} + assert Nx.shape(pooled) == {1, 2} + + # Verify L2 normalization: norm should equal 1.0 + norm = embedding |> Nx.LinAlg.norm() |> Nx.to_number() + assert_in_delta norm, 1.0, 1.0e-5 + end + + @tag :tmp_dir + test "supports fuse_dense option to merge consecutive linear layers", %{ + model_info: model_info, + tmp_dir: tmp_dir + } do + modules = [ + %{ + "idx" => 0, + "name" => "0", + "path" => "", + "type" => "sentence_transformers.models.Transformer" + }, + %{ + "idx" => 1, + "name" => "1", + "path" => "1_Dense", + "type" => "sentence_transformers.models.Dense" + }, + %{ + "idx" => 2, + "name" => "2", + "path" => "2_Dense", + "type" => "sentence_transformers.models.Dense" + } + ] + + File.write!(Path.join(tmp_dir, "modules.json"), Jason.encode!(modules)) + + # 1_Dense: 4 -> 8 (no bias, identity activation) + File.mkdir_p!(Path.join(tmp_dir, "1_Dense")) + + c1 = %{ + "in_features" => 4, + "out_features" => 8, + "bias" => false, + "activation_function" => nil + } + + File.write!(Path.join(tmp_dir, "1_Dense/config.json"), Jason.encode!(c1)) + w1 = Nx.broadcast(0.5, {8, 4}) + + Safetensors.write!(Path.join(tmp_dir, "1_Dense/model.safetensors"), %{"linear.weight" => w1}) + + # 2_Dense: 8 -> 2 (no bias, identity activation) + File.mkdir_p!(Path.join(tmp_dir, "2_Dense")) + + c2 = %{ + "in_features" => 8, + "out_features" => 2, + "bias" => false, + "activation_function" => nil + } + + File.write!(Path.join(tmp_dir, "2_Dense/config.json"), Jason.encode!(c2)) + w2 = Nx.broadcast(0.25, {2, 8}) + + Safetensors.write!(Path.join(tmp_dir, "2_Dense/model.safetensors"), %{"linear.weight" => w2}) + + assert {:ok, fused_info} = + SentenceTransformers.load_embedding_head({:local, tmp_dir}, model_info, + fuse_dense: true + ) + + assert Map.has_key?(fused_info.params.data, "1_Dense_2_Dense") + fused_kernel = fused_info.params.data["1_Dense_2_Dense"]["kernel"] + assert Nx.shape(fused_kernel) == {4, 2} + 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 = 0.5 |> Nx.broadcast({8, 4}) |> Nx.as_type({:f, 32}) + b1 = 0.1 |> Nx.broadcast({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 = 0.25 |> Nx.broadcast({4, 8}) |> Nx.as_type({:f, 32}) + b2 = 0.2 |> Nx.broadcast({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} = + SentenceTransformers.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 "loads max_tokens pooling", %{tmp_dir: dir} do + custom_model = + "input_ids" + |> Axon.input(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{data: %{}, parameters: %{}}, + spec: nil + } + + 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} = + SentenceTransformers.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 with both key variants", %{tmp_dir: dir} do + custom_model = + "input_ids" + |> Axon.input(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{data: %{}, parameters: %{}}, + spec: nil + } + + 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} = + SentenceTransformers.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", %{tmp_dir: dir} do + custom_model = + "input_ids" + |> Axon.input(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{data: %{}, parameters: %{}}, + spec: nil + } + + 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} = + SentenceTransformers.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] + 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 "loads weightedmean_tokens pooling", %{tmp_dir: dir} do + custom_model = + "input_ids" + |> Axon.input(shape: {nil, nil}) + |> Axon.nx(fn _input_ids -> + hidden_state = Nx.tensor([[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]]) + %{hidden_state: hidden_state} + end) + + model_info = %{ + model: custom_model, + params: %Axon.ModelState{data: %{}, parameters: %{}}, + spec: nil + } + + 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_weightedmean_tokens" => true + } + + File.write!(Path.join(pooling_dir, "config.json"), Jason.encode!(pooling_config)) + + assert {:ok, new_model_info} = + SentenceTransformers.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) + + # 1 * [1, 2] + 2 * [3, 4] = [7, 10]; sum weights = 3 + assert_all_close(output.embedding, Nx.tensor([[7.0 / 3, 10.0 / 3]])) + end + + @tag :tmp_dir + test "excludes prompt tokens from pooling when include_prompt is false", %{ + tmp_dir: dir + } do + custom_model = + "input_ids" + |> Axon.input(shape: {nil, nil}) + |> Axon.nx(fn _input_ids -> + hidden_state = + Nx.tensor([ + [[1.0, 0.0], [2.0, 1.0], [3.0, 2.0], [4.0, 3.0]], + [[5.0, 4.0], [6.0, 5.0], [7.0, 6.0], [8.0, 7.0]] + ]) + + %{hidden_state: hidden_state} + end) + + model_info = %{ + model: custom_model, + params: %Axon.ModelState{data: %{}, parameters: %{}}, + spec: nil + } + + 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_mean_tokens" => true, + "include_prompt" => false + } + + File.write!(Path.join(pooling_dir, "config.json"), Jason.encode!(pooling_config)) + + assert {:ok, new_model_info} = + SentenceTransformers.load_embedding_head({:local, dir}, model_info) + + inputs = %{ + "input_ids" => Nx.tensor([[1, 2, 3, 4], [1, 2, 3, 4]]), + "attention_mask" => Nx.tensor([[1, 1, 1, 1], [1, 1, 1, 1]]), + "prompt_length" => Nx.tensor([1, 2]) + } + + {_init_fn, predict_fn} = Axon.build(new_model_info.model) + output = predict_fn.(new_model_info.params, inputs) + + # First sequence: exclude 1 prompt token -> mean of [[2,1], [3,2], [4,3]] = [3, 2] + # Second sequence: exclude 2 prompt tokens -> mean of [[7,6], [8,7]] = [7.5, 6.5] + assert_all_close(output.embedding, Nx.tensor([[3.0, 2.0], [7.5, 6.5]])) + end + + @tag :tmp_dir + test "excludes prompt tokens with left padding when include_prompt is false", %{ + tmp_dir: dir + } do + custom_model = + "input_ids" + |> Axon.input(shape: {nil, nil}) + |> Axon.nx(fn _input_ids -> + hidden_state = + Nx.tensor([ + [[0.0, 0.0], [0.0, 0.0], [1.0, 0.0], [2.0, 1.0], [3.0, 2.0], [4.0, 3.0]] + ]) + + %{hidden_state: hidden_state} + end) + + model_info = %{ + model: custom_model, + params: %Axon.ModelState{data: %{}, parameters: %{}}, + spec: nil + } + + 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_mean_tokens" => true, + "include_prompt" => false + } + + File.write!(Path.join(pooling_dir, "config.json"), Jason.encode!(pooling_config)) + + assert {:ok, new_model_info} = + SentenceTransformers.load_embedding_head({:local, dir}, model_info) + + inputs = %{ + "input_ids" => Nx.tensor([[1, 2, 3, 4, 5, 6]]), + "attention_mask" => Nx.tensor([[0, 0, 1, 1, 1, 1]]), + "prompt_length" => Nx.tensor([1]) + } + + {_init_fn, predict_fn} = Axon.build(new_model_info.model) + output = predict_fn.(new_model_info.params, inputs) + + # 2 pad tokens + 1 prompt token excluded -> mean of [[2,1], [3,2], [4,3]] = [3, 2] + assert_all_close(output.embedding, Nx.tensor([[3.0, 2.0]])) + end + + @tag :tmp_dir + test "loads LayerNorm module", %{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_LayerNorm", "type" => "sentence_transformers.models.LayerNorm"} + ] + + 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}) + ) + + ln_dir = Path.join(dir, "2_LayerNorm") + File.mkdir_p!(ln_dir) + + File.write!(Path.join(ln_dir, "config.json"), Jason.encode!(%{"eps" => 1.0e-5})) + + tensors = %{ + "weight" => Nx.tensor([2.0, 2.0, 2.0, 2.0]), + "bias" => Nx.tensor([1.0, 1.0, 1.0, 1.0]) + } + + Safetensors.write!(Path.join(ln_dir, "model.safetensors"), tensors) + + assert {:ok, new_model_info} = + SentenceTransformers.load_embedding_head({:local, dir}, model_info) + + inputs = %{ + "input_ids" => Nx.tensor([[1, 2]]), + "attention_mask" => Nx.tensor([[1, 1]]) + } + + {_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([[1.0, 1.0, 1.0, 1.0]])) + 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) + + 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" => 0.5 |> Nx.broadcast({4, 4}) |> Nx.as_type({:f, 32}) + } + + Safetensors.write!(Path.join(dense_dir, "model.safetensors"), weights) + + assert {:ok, new_model_info} = + SentenceTransformers.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) + end + + @tag :tmp_dir + test "loads config metadata with fallback filenames", %{tmp_dir: dir} do + # Test fallback to sentence_bert_config.json + st_config = %{ + "prompts" => %{ + "query" => "Represent this query: ", + "passage" => "Represent this passage: " + }, + "default_prompt_name" => "query", + "similarity_fn_name" => "cosine", + "max_seq_length" => 256 + } + + File.write!(Path.join(dir, "sentence_bert_config.json"), Jason.encode!(st_config)) + + assert {:ok, config} = SentenceTransformers.load_config({:local, dir}) + assert config.default_prompt_name == "query" + assert config.similarity_fn_name == :cosine + assert config.max_seq_length == 256 + assert config.prompts["query"] == "Represent this query: " + end + + @tag :tmp_dir + test "loads config metadata from 0_Transformer/sentence_bert_config.json fallback", %{ + tmp_dir: dir + } do + transformer_dir = Path.join(dir, "0_Transformer") + File.mkdir_p!(transformer_dir) + + st_config = %{ + "prompts" => %{"query" => "Represent query from 0_Transformer: "}, + "default_prompt_name" => "query", + "similarity_fn_name" => "dot" + } + + File.write!( + Path.join(transformer_dir, "sentence_bert_config.json"), + Jason.encode!(st_config) + ) + + assert {:ok, config} = SentenceTransformers.load_config({:local, dir}) + assert config.default_prompt_name == "query" + assert config.similarity_fn_name == :dot + assert config.prompts["query"] == "Represent query from 0_Transformer: " + end + + @tag :tmp_dir + test "config priority: config_sentence_transformers.json > sentence_bert_config.json > 0_Transformer/sentence_bert_config.json", + %{tmp_dir: dir} do + transformer_dir = Path.join(dir, "0_Transformer") + File.mkdir_p!(transformer_dir) + + File.write!( + Path.join(transformer_dir, "sentence_bert_config.json"), + Jason.encode!(%{"default_prompt_name" => "from_0_transformer"}) + ) + + File.write!( + Path.join(dir, "sentence_bert_config.json"), + Jason.encode!(%{"default_prompt_name" => "from_root_legacy"}) + ) + + # 1. sentence_bert_config.json takes precedence over 0_Transformer + assert {:ok, config} = SentenceTransformers.load_config({:local, dir}) + assert config.default_prompt_name == "from_root_legacy" + + # 2. config_sentence_transformers.json takes precedence over both + File.write!( + Path.join(dir, "config_sentence_transformers.json"), + Jason.encode!(%{"default_prompt_name" => "from_canonical"}) + ) + + assert {:ok, config} = SentenceTransformers.load_config({:local, dir}) + assert config.default_prompt_name == "from_canonical" + end + + test "causal fallback: infers 'last' pooling for Bumblebee causal specs" do + causal_specs = [ + struct(Bumblebee.Text.Gemma), + struct(Bumblebee.Text.Gemma3Text), + struct(Bumblebee.Text.Llama), + struct(Bumblebee.Text.Mistral), + struct(Bumblebee.Text.Phi), + struct(Bumblebee.Text.Phi3), + struct(Bumblebee.Text.Gpt2), + struct(Bumblebee.Text.GptBigCode), + struct(Bumblebee.Text.GptNeoX), + struct(Bumblebee.Text.Qwen3), + struct(Bumblebee.Text.SmolLm3) + ] + + for spec <- causal_specs do + model_info = %{spec: spec} + + assert SentenceTransformers.infer_fallback_pooling_mode( + {:local, "nonexistent"}, + model_info + ) == "last" + end + end + + test "infers 'mean' pooling for Bumblebee non-causal specs" do + non_causal_specs = [ + struct(Bumblebee.Text.Bert), + struct(Bumblebee.Text.Roberta), + struct(Bumblebee.Text.Distilbert), + struct(Bumblebee.Text.Albert) + ] + + for spec <- non_causal_specs do + model_info = %{spec: spec} + + assert SentenceTransformers.infer_fallback_pooling_mode( + {:local, "nonexistent"}, + model_info + ) == "mean" + end + end + + @tag :tmp_dir + test "infers 'last' pooling for config.json with *ForCausalLM or *LMHeadModel architectures", + %{tmp_dir: dir} do + causal_archs = [ + ["LlamaForCausalLM"], + ["MistralForCausalLM"], + ["GemmaForCausalLM"], + ["PhiForCausalLM"], + ["Qwen2ForCausalLM"], + ["GPT2LMHeadModel"] + ] + + for archs <- causal_archs do + File.write!(Path.join(dir, "config.json"), Jason.encode!(%{"architectures" => archs})) + + assert SentenceTransformers.infer_fallback_pooling_mode({:local, dir}, %{spec: nil}) == + "last" + end + end + + @tag :tmp_dir + test "infers 'mean' pooling for non-causal architectures or when is_causal is false", %{ + tmp_dir: dir + } do + File.write!( + Path.join(dir, "config.json"), + Jason.encode!(%{"architectures" => ["BertModel"]}) + ) + + assert SentenceTransformers.infer_fallback_pooling_mode({:local, dir}, %{spec: nil}) == + "mean" + + File.write!( + Path.join(dir, "config.json"), + Jason.encode!(%{ + "architectures" => ["LlamaForCausalLM"], + "is_causal" => false + }) + ) + + assert SentenceTransformers.infer_fallback_pooling_mode({:local, dir}, %{spec: nil}) == + "mean" + 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"} = + SentenceTransformers.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 Transformer or StaticEmbedding"} = + SentenceTransformers.load_embedding_head({:local, dir}, model_info) + 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 \"Unknown\""} = + SentenceTransformers.load_embedding_head({:local, dir}, model_info) + end + + @tag :tmp_dir + test "returns error on unsupported activation function", %{ + 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.UnsupportedAct" + }) + ) + + assert {:error, + "unsupported activation function \"torch.nn.modules.activation.UnsupportedAct\" in 2_Dense"} = + SentenceTransformers.load_embedding_head({:local, dir}, model_info) + end + + @tag :tmp_dir + test "loads dense layer with Sigmoid activation", %{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" + }) + ) + + weights = %{ + "linear.weight" => 0.0 |> Nx.broadcast({4, 4}) |> Nx.as_type({:f, 32}) + } + + Safetensors.write!(Path.join(dense_dir, "model.safetensors"), weights) + + assert {:ok, new_model_info} = + SentenceTransformers.load_embedding_head({:local, dir}, model_info) + + inputs = %{ + "input_ids" => Nx.tensor([[1, 2]]), + "attention_mask" => Nx.tensor([[1, 1]]) + } + + {_init_fn, predict_fn} = Axon.build(new_model_info.model) + output = predict_fn.(new_model_info.params, inputs) + + # sigmoid(0) = 0.5 + assert_all_close(output.embedding, Nx.tensor([[0.5, 0.5, 0.5, 0.5]])) + end + + @tag :tmp_dir + test "loads dense layer with use_residual: 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"} + ] + + 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" => nil, + "use_residual" => true + }) + ) + + weights = %{ + "linear.weight" => 0.5 |> Nx.broadcast({4, 4}) |> Nx.as_type({:f, 32}) + } + + Safetensors.write!(Path.join(dense_dir, "model.safetensors"), weights) + + assert {:ok, new_model_info} = + SentenceTransformers.load_embedding_head({:local, dir}, model_info) + + inputs = %{ + "input_ids" => Nx.tensor([[1, 2]]), + "attention_mask" => Nx.tensor([[1, 1]]) + } + + {_init_fn, predict_fn} = Axon.build(new_model_info.model) + output = predict_fn.(new_model_info.params, inputs) + + # input pooled is [1, 1, 1, 1]. linear(x) = [2, 2, 2, 2]. with residual: [3, 3, 3, 3] + assert_all_close(output.embedding, Nx.tensor([[3.0, 3.0, 3.0, 3.0]])) + end + + @tag :tmp_dir + test "supports truncate_dim and re-normalization", %{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) + + File.write!( + Path.join(pooling_dir, "config.json"), + Jason.encode!(%{"pooling_mode_mean_tokens" => true}) + ) + + assert {:ok, new_model_info} = + SentenceTransformers.load_embedding_head({:local, dir}, model_info, + truncate_dim: 2 + ) + + inputs = %{ + "input_ids" => Nx.tensor([[1, 2]]), + "attention_mask" => Nx.tensor([[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, 2} + # Re-normalized L2 norm must be 1.0 + norm = output.embedding |> Nx.LinAlg.norm() |> Nx.to_number() + assert_in_delta norm, 1.0, 1.0e-5 + end + + test "truncate_embeddings helper function" do + emb = Nx.tensor([[1.0, 1.0, 1.0, 1.0]]) + truncated = SentenceTransformers.truncate_embeddings(emb, 2) + assert Nx.shape(truncated) == {1, 2} + assert_all_close(truncated, Nx.tensor([[1.0, 1.0]])) + + truncated_norm = SentenceTransformers.truncate_embeddings(emb, 2, normalize: true) + norm = truncated_norm |> Nx.LinAlg.norm() |> Nx.to_number() + assert_in_delta norm, 1.0, 1.0e-5 + end + + @tag :tmp_dir + test "loads WeightedLayerPooling module", %{tmp_dir: dir} do + custom_model = + "input_ids" + |> Axon.input(shape: {nil, nil}) + |> Axon.nx(fn _input_ids -> + l1 = Nx.broadcast(1.0, {1, 2, 4}) + l2 = Nx.broadcast(2.0, {1, 2, 4}) + l3 = Nx.broadcast(3.0, {1, 2, 4}) + stacked = Nx.stack([l1, l2, l3], axis: 0) + %{hidden_state: stacked} + end) + + model_info = %{ + model: custom_model, + params: %Axon.ModelState{data: %{}, parameters: %{}}, + spec: nil + } + + modules = [ + %{"idx" => 0, "path" => "", "type" => "sentence_transformers.models.Transformer"}, + %{ + "idx" => 1, + "path" => "1_WeightedLayerPooling", + "type" => "sentence_transformers.models.WeightedLayerPooling" + }, + %{"idx" => 2, "path" => "2_Pooling", "type" => "sentence_transformers.models.Pooling"} + ] + + File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules)) + + wlp_dir = Path.join(dir, "1_WeightedLayerPooling") + File.mkdir_p!(wlp_dir) + + File.write!( + Path.join(wlp_dir, "config.json"), + Jason.encode!(%{"layer_start" => 0, "num_hidden_layers" => 2}) + ) + + weights = %{ + "layer_weights" => Nx.tensor([1.0, 2.0, 3.0], type: {:f, 32}) + } + + Safetensors.write!(Path.join(wlp_dir, "model.safetensors"), weights) + + pooling_dir = Path.join(dir, "2_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} = + SentenceTransformers.load_embedding_head({:local, dir}, model_info) + + inputs = %{ + "input_ids" => Nx.tensor([[1, 2]]), + "attention_mask" => Nx.tensor([[1, 1]]) + } + + {_init_fn, predict_fn} = Axon.build(new_model_info.model) + output = predict_fn.(new_model_info.params, inputs) + + expected = Nx.broadcast(2.3333333, {1, 4}) + assert_all_close(output.embedding, expected) + end + + @tag :tmp_dir + test "loads StaticEmbedding model (Model2Vec style)", %{tmp_dir: dir} do + modules = [ + %{ + "idx" => 0, + "name" => "0", + "path" => "0_StaticEmbedding", + "type" => "sentence_transformers.models.StaticEmbedding" + }, + %{ + "idx" => 1, + "name" => "1", + "path" => "1_Pooling", + "type" => "sentence_transformers.models.Pooling" + } + ] + + File.write!(Path.join(dir, "modules.json"), Jason.encode!(modules)) + + static_dir = Path.join(dir, "0_StaticEmbedding") + File.mkdir_p!(static_dir) + + emb_matrix = Nx.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], type: {:f, 32}) + + Safetensors.write!(Path.join(static_dir, "model.safetensors"), %{ + "embedding.weight" => emb_matrix + }) + + 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, model_info} = SentenceTransformers.load_model({:local, dir}) + + inputs = %{ + "input_ids" => Nx.tensor([[0, 2]]), + "attention_mask" => Nx.tensor([[1, 1]]) + } + + {_init_fn, predict_fn} = Axon.build(model_info.model) + output = predict_fn.(model_info.params, inputs) + + assert_all_close(output.embedding, Nx.tensor([[3.0, 4.0]])) + end + + @tag :tmp_dir + test "updates ModelState.parameters for added layers", %{tmp_dir: dir} do + base_model = + "input_ids" + |> Axon.input(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: nil} + + 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" => 0.5 |> Nx.broadcast({8, 4}) |> Nx.as_type({:f, 32}), + "linear.bias" => 0.1 |> Nx.broadcast({8}) |> Nx.as_type({:f, 32}) + } + + Safetensors.write!(Path.join(dense_dir, "model.safetensors"), weights) + + assert {:ok, new_model_info} = + SentenceTransformers.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 + end + + @tag :slow + @tag :network + test "works directly with Bumblebee.Text.text_embedding and piping" do + repo = {:hf, "sentence-transformers/all-MiniLM-L6-v2"} + + # Piping: Bumblebee.load_model |> SentenceTransformers.load_embedding_head + assert {:ok, model_info} = + repo + |> Bumblebee.load_model(architecture: :base) + |> SentenceTransformers.load_embedding_head(repo) + + assert {:ok, tokenizer} = Bumblebee.load_tokenizer(repo) + + # Uses standard Bumblebee serving directly! + serving = Bumblebee.Text.text_embedding(model_info, tokenizer) + result = Nx.Serving.run(serving, "Hello world") + + assert %{embedding: embedding} = result + assert Nx.shape(embedding) == {384} + + first_5 = Nx.to_flat_list(embedding[0..4]) + py_reference = [-0.03447720, 0.03102321, 0.00673497, 0.02610898, -0.03936199] + + first_5 + |> Enum.zip(py_reference) + |> Enum.each(fn {val, ref} -> + assert_in_delta val, ref, 1.0e-5 + end) + end + + @tag :slow + @tag :network + test "end-to-end integration test with sentence-transformers/all-MiniLM-L6-v2" do + repo = {:hf, "sentence-transformers/all-MiniLM-L6-v2"} + + assert {:ok, model_info} = SentenceTransformers.load_model(repo) + assert {:ok, tokenizer} = Bumblebee.load_tokenizer(repo) + + serving = Bumblebee.Text.text_embedding(model_info, tokenizer) + result = Nx.Serving.run(serving, "Hello world") + + assert %{embedding: embedding} = result + assert Nx.shape(embedding) == {384} + + # Reference values from PyTorch sentence-transformers 6.1.0 for "Hello world": + # [-0.034477, 0.031023, 0.006735, 0.026109, -0.039362] + first_5 = Nx.to_flat_list(embedding[0..4]) + py_reference = [-0.03447720, 0.03102321, 0.00673497, 0.02610898, -0.03936199] + + first_5 + |> Enum.zip(py_reference) + |> Enum.each(fn {val, ref} -> + assert_in_delta val, ref, 1.0e-5 + end) + end +end diff --git a/test/support/sentence_transformers/python_bridge.ex b/test/support/sentence_transformers/python_bridge.ex new file mode 100644 index 00000000..aea025db --- /dev/null +++ b/test/support/sentence_transformers/python_bridge.ex @@ -0,0 +1,557 @@ +defmodule Bumblebee.SentenceTransformers.TestSupport.PythonBridge do + @moduledoc """ + Helper module for executing Python inference scripts (via `uv` or `python3`) + and generating random SentenceTransformers / Hugging Face model checkpoints + to compare outputs with Bumblebee.SentenceTransformers. + """ + + @doc """ + Runs inference using Python's `sentence_transformers` library on a given model and text, + returning the computed embeddings as an `Nx.Tensor`. + """ + def run_sentence_transformers(repo_id, text, opts \\ []) do + code = """ + import json + import os + import sys + import warnings + warnings.filterwarnings("ignore") + + os.environ["TOKENIZERS_PARALLELISM"] = "false" + os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" + + from sentence_transformers import SentenceTransformer + + model = SentenceTransformer("#{repo_id}") + encode_opts = json.loads(#{inspect(Jason.encode!(Map.new(Keyword.take(opts, [:prompt, :prompt_name]))))}) + res = model.encode(#{inspect(text)}, **encode_opts) + print("__RESULT__:" + json.dumps(res.tolist())) + """ + + deps = ["sentence-transformers", "torch"] + run_python(code, deps, opts) + end + + @pooling_mode_flags %{ + "mean_tokens" => "pooling_mode_mean_tokens=True", + "cls_token" => "pooling_mode_cls_token=True", + "max_tokens" => "pooling_mode_max_tokens=True", + "mean_sqrt_len_tokens" => "pooling_mode_mean_sqrt_len_tokens=True", + "weightedmean_tokens" => "pooling_mode_weightedmean_tokens=True", + "last_token" => "pooling_mode_lasttoken=True" + } + + @doc """ + Generates an on-the-fly random SentenceTransformer model checkpoint in Python with random weights, + saves it to a temporary directory, executes inference on `text`, and returns: + `{:ok, %{dir: dir, python_output: nx_tensor, text: text}}`. + + ## Options + * `:pooling_mode` - a single pooling mode name or a list of mode names. + Supported: `"mean_tokens"`, `"cls_token"`, `"max_tokens"`, + `"mean_sqrt_len_tokens"`, `"weightedmean_tokens"`, `"last_token"` + (default: `["mean_tokens"]`). + * `:with_dense` - whether to include a Dense projection layer (default: `false`) + * `:dense_activation` - activation for the Dense layer: `"Identity"`, `"Tanh"`, + `"ReLU"`, `"GELU"`, `"SiLU"`, `"Sigmoid"`, `"Mish"`, `"LeakyReLU"` + (default: `"Tanh"`) + * `:with_normalize` - whether to append a Normalize module (default: `false`) + * `:with_layernorm` - whether to append a LayerNorm module (default: `false`) + * `:with_dropout` - whether to append a Dropout module (default: `false`) + * `:with_fused_dense` - whether to append two consecutive Dense layers + (default: `false`). By default both use `Identity` activation and no bias so + Bumblebee.SentenceTransformers can fuse them. + * `:fused_dense_first_bias` - whether the first fused Dense layer has bias + (default: `false`) + * `:fused_dense_first_activation` - activation for the first fused Dense layer + (default: `"Identity"`) + * `:fused_dense_second_bias` - whether the second fused Dense layer has bias + (default: `false`) + * `:fused_dense_second_activation` - activation for the second fused Dense layer + (default: `"Identity"`) + * `:prompt` - prompt string to prepend to the input (default: `nil`) + * `:prompt_name` - name of a prompt in `config_sentence_transformers.json` to + use for the Python encode call (default: `nil`) + * `:document_prompt` - prompt string to register under the `"document"` name + (default: `nil`) + * `:include_prompt` - whether to include prompt tokens in pooling + (default: `true`). Only relevant when `:prompt` is given. + * `:text` - text to encode (default: `"Hello world"`) + """ + def generate_random_sentence_transformer(opts \\ []) do + pooling_mode = Keyword.get(opts, :pooling_mode, ["mean_tokens"]) + with_dense = Keyword.get(opts, :with_dense, false) + dense_activation = Keyword.get(opts, :dense_activation, "Tanh") + with_normalize = Keyword.get(opts, :with_normalize, false) + with_layernorm = Keyword.get(opts, :with_layernorm, false) + with_dropout = Keyword.get(opts, :with_dropout, false) + with_fused_dense = Keyword.get(opts, :with_fused_dense, false) + fused_dense_first_bias = Keyword.get(opts, :fused_dense_first_bias, false) + fused_dense_first_activation = Keyword.get(opts, :fused_dense_first_activation, "Identity") + fused_dense_second_bias = Keyword.get(opts, :fused_dense_second_bias, false) + fused_dense_second_activation = Keyword.get(opts, :fused_dense_second_activation, "Identity") + prompt = Keyword.get(opts, :prompt) + include_prompt = Keyword.get(opts, :include_prompt, true) + text = Keyword.get(opts, :text, "Hello world") + max_seq_length = Keyword.get(opts, :max_seq_length, 64) + + mode_flags = + pooling_mode + |> List.wrap() + |> Enum.map_join(", ", &Map.fetch!(@pooling_mode_flags, &1)) + + mode_flags = + if "mean_tokens" in List.wrap(pooling_mode), + do: mode_flags, + else: "pooling_mode_mean_tokens=False, " <> mode_flags + + pooling_args = + if prompt do + "#{mode_flags}, include_prompt=#{if include_prompt, do: "True", else: "False"}" + else + mode_flags + end + + prompt_name = Keyword.get(opts, :prompt_name) + + encode_call = + cond do + prompt_name -> + "st_model.encode(#{inspect(text)}, prompt_name=#{inspect(prompt_name)})" + + prompt -> + "st_model.encode(#{inspect(text)}, prompt=#{inspect(prompt)})" + + true -> + "st_model.encode(#{inspect(text)})" + end + + prompts = + if prompt do + %{"query" => prompt} + else + %{} + end + + prompts = + if document_prompt = Keyword.get(opts, :document_prompt) do + Map.put(prompts, "document", document_prompt) + else + prompts + end + + st_config = + if map_size(prompts) > 0 do + Jason.encode!(%{ + "prompts" => prompts, + "default_prompt_name" => "query" + }) + else + Jason.encode!(%{}) + end + + code = """ + import json + import os + import sys + import tempfile + import warnings + warnings.filterwarnings("ignore") + + os.environ["TOKENIZERS_PARALLELISM"] = "false" + os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" + + import torch + from transformers import BertConfig, BertModel, BertTokenizerFast + from sentence_transformers import SentenceTransformer, models + + torch.manual_seed(42) + + hidden_size = 32 + config = BertConfig( + vocab_size=128, + hidden_size=hidden_size, + num_hidden_layers=2, + num_attention_heads=4, + intermediate_size=64, + max_position_embeddings=64 + ) + base_model = BertModel(config) + + with tempfile.TemporaryDirectory() as tmp_base: + base_model.save_pretrained(tmp_base) + vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]", "hello", "world", + "represent", "the", "sentence", "for", "retrieval", ":", "query", "document"] + vocab += [f"unused_{i}" for i in range(128 - len(vocab))] + with open(os.path.join(tmp_base, "vocab.txt"), "w") as f: + f.write("\\n".join(vocab)) + BertTokenizerFast(vocab={token: i for i, token in enumerate(vocab)}, do_lower_case=True).save_pretrained(tmp_base) + + word_embedding_model = models.Transformer(tmp_base, max_seq_length=#{max_seq_length}) + pooling_model = models.Pooling( + word_embedding_model.get_embedding_dimension(), + #{pooling_args} + ) + + modules = [word_embedding_model, pooling_model] + + activation_map = { + "Identity": torch.nn.Identity(), + "Tanh": torch.nn.Tanh(), + "ReLU": torch.nn.ReLU(), + "GELU": torch.nn.GELU(), + "SiLU": torch.nn.SiLU(), + "Sigmoid": torch.nn.Sigmoid(), + "Mish": torch.nn.Mish(), + "LeakyReLU": torch.nn.LeakyReLU() + } + + if #{if with_dense, do: "True", else: "False"}: + dense_model = models.Dense( + in_features=hidden_size, + out_features=16, + bias=True, + activation_function=activation_map[#{inspect(dense_activation)}] + ) + modules.append(dense_model) + + if #{if with_fused_dense, do: "True", else: "False"}: + modules.append(models.Dense( + in_features=hidden_size, + out_features=24, + bias=#{if fused_dense_first_bias, do: "True", else: "False"}, + activation_function=activation_map[#{inspect(fused_dense_first_activation)}] + )) + modules.append(models.Dense( + in_features=24, + out_features=16, + bias=#{if fused_dense_second_bias, do: "True", else: "False"}, + activation_function=activation_map[#{inspect(fused_dense_second_activation)}] + )) + + if #{if with_normalize, do: "True", else: "False"}: + modules.append(models.Normalize()) + + if #{if with_layernorm, do: "True", else: "False"}: + modules.append(models.LayerNorm(hidden_size)) + + if #{if with_dropout, do: "True", else: "False"}: + modules.append(models.Dropout(0.5)) + + st_model = SentenceTransformer(modules=modules) + + save_dir = tempfile.mkdtemp(prefix="st_random_st_") + st_model.save(save_dir) + + with open(save_dir + "/config_sentence_transformers.json", "w") as f: + f.write(#{inspect(st_config)}) + + st_model = SentenceTransformer(save_dir) + + emb = #{encode_call}.tolist() + + print("__DIR__:" + save_dir) + print("__RESULT__:" + json.dumps(emb)) + """ + + deps = ["sentence-transformers", "torch"] + {cmd, args} = python_exec_args(code, deps) + + case System.cmd(cmd, args, stderr_to_stdout: true) do + {output, 0} -> + [_, dir] = Regex.run(~r/__DIR__:(.*)$/m, output) + [_, result_json] = Regex.run(~r/__RESULT__:(.*)$/m, output) + + python_output = Nx.tensor(Jason.decode!(result_json), type: :f32) + {:ok, %{dir: String.trim(dir), python_output: python_output, text: text}} + + {output, exit_code} -> + raise "Failed to generate random SentenceTransformer (exit #{exit_code}):\n#{output}" + end + end + + @doc """ + Runs inference using Python's `transformers` library on a given model and inputs. + """ + def run_hf_model(repo_id, inputs, opts \\ []) do + output_key = Keyword.get(opts, :output_key, "last_hidden_state") + json_inputs = Jason.encode!(inputs) + + code = """ + import json + import os + import sys + import warnings + warnings.filterwarnings("ignore") + + os.environ["TOKENIZERS_PARALLELISM"] = "false" + os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" + + import torch + from transformers import AutoModel + + raw_inputs = json.loads(#{inspect(json_inputs)}) + inputs = {k: torch.tensor(v) for k, v in raw_inputs.items()} + + model = AutoModel.from_pretrained("#{repo_id}") + model.eval() + + with torch.no_grad(): + outputs = model(**inputs) + result = getattr(outputs, "#{output_key}", None) + if result is None and "#{output_key}" in outputs: + result = outputs["#{output_key}"] + if hasattr(result, "cpu"): + result = result.cpu().numpy().tolist() + + print("__RESULT__:" + json.dumps(result)) + """ + + deps = ["transformers", "torch"] + run_python(code, deps, opts) + end + + @doc """ + Generates a tiny model with completely random weights on the fly in Python, saves it to + a temporary directory, executes inference in Python, and returns: + `%{dir: temp_dir, inputs: inputs_map, python_output: nx_tensor}`. + """ + def generate_random_hf_model(model_family, opts \\ []) do + seq_len = Keyword.get(opts, :seq_len, 5) + batch_size = Keyword.get(opts, :batch_size, 1) + + code = """ + import json + import os + import sys + import tempfile + import warnings + warnings.filterwarnings("ignore") + + os.environ["TOKENIZERS_PARALLELISM"] = "false" + os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" + + import torch + from transformers import BertConfig, BertModel, RobertaConfig, RobertaModel, GPT2Config, GPT2Model + + family = "#{model_family}" + batch_size = #{batch_size} + seq_len = #{seq_len} + + torch.manual_seed(42) + + if family == "bert": + config = BertConfig( + vocab_size=128, + hidden_size=32, + num_hidden_layers=2, + num_attention_heads=4, + intermediate_size=64, + max_position_embeddings=64 + ) + model = BertModel(config) + inputs = { + "input_ids": torch.randint(0, 128, (batch_size, seq_len)), + "attention_mask": torch.ones((batch_size, seq_len), dtype=torch.long) + } + output_attr = "last_hidden_state" + + elif family == "roberta": + config = RobertaConfig( + vocab_size=128, + hidden_size=32, + num_hidden_layers=2, + num_attention_heads=4, + intermediate_size=64, + max_position_embeddings=64, + type_vocab_size=1 + ) + model = RobertaModel(config) + inputs = { + "input_ids": torch.randint(0, 128, (batch_size, seq_len)), + "attention_mask": torch.ones((batch_size, seq_len), dtype=torch.long) + } + output_attr = "last_hidden_state" + + elif family == "gpt2": + config = GPT2Config( + vocab_size=128, + n_embd=32, + n_layer=2, + n_head=4, + n_positions=64, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2 + ) + model = GPT2Model(config) + inputs = { + "input_ids": torch.randint(0, 128, (batch_size, seq_len)), + "attention_mask": torch.ones((batch_size, seq_len), dtype=torch.long) + } + output_attr = "last_hidden_state" + else: + raise ValueError(f"Unsupported family: {family}") + + model.eval() + + with torch.no_grad(): + outputs = model(**inputs) + result = getattr(outputs, output_attr).cpu().numpy().tolist() + + temp_dir = tempfile.mkdtemp(prefix="st_random_hf_") + model.save_pretrained(temp_dir) + + raw_inputs = {k: v.tolist() for k, v in inputs.items()} + + print("__DIR__:" + temp_dir) + print("__INPUTS__:" + json.dumps(raw_inputs)) + print("__RESULT__:" + json.dumps(result)) + """ + + deps = ["transformers", "torch"] + {cmd, args} = python_exec_args(code, deps) + + case System.cmd(cmd, args, stderr_to_stdout: true) do + {output, 0} -> + [_, dir] = Regex.run(~r/__DIR__:(.*)$/m, output) + [_, result_json] = Regex.run(~r/__RESULT__:(.*)$/m, output) + + inputs = + Regex.run(~r/__INPUTS__:(.*)$/m, output) + |> Enum.at(1) + |> Jason.decode!() + |> Map.new(fn {k, v} -> {k, Nx.tensor(v)} end) + + python_output = Nx.tensor(Jason.decode!(result_json), type: :f32) + + {:ok, %{dir: String.trim(dir), inputs: inputs, python_output: python_output}} + + {output, exit_code} -> + raise "Failed to generate random model (exit #{exit_code}):\n#{output}" + end + end + + @doc """ + Runs custom Python code that prints `__RESULT__:` and returns the parsed Nx tensor or raw data. + """ + def run_python_code(code, deps \\ ["sentence-transformers", "torch"], opts \\ []) do + {cmd, args} = python_exec_args(code, deps) + + case System.cmd(cmd, args, stderr_to_stdout: true) do + {output, 0} -> + parse_result(output, opts) + + {output, exit_code} -> + raise "Python process failed with exit code #{exit_code}:\n#{output}" + end + end + + defp run_python(code, deps, opts) do + run_python_code(code, deps, opts) + end + + defp python_exec_args(code, deps) do + {cmd, args} = python_command(deps) + verify_environment!(cmd, args, deps) + {cmd, args ++ ["-c", code]} + end + + defp verify_environment!(cmd, args, deps) do + key = {__MODULE__, :environment, cmd, args, deps} + + result = + case :persistent_term.get(key, :unchecked) do + :unchecked -> + code = """ + import importlib.metadata as metadata + import json, platform, sys + required = #{Jason.encode!(deps)} + versions = {} + try: + for package in required: + versions[package] = metadata.version(package) + if "sentence-transformers" in versions and versions["sentence-transformers"] != "6.1.0": + raise RuntimeError("expected sentence-transformers==6.1.0, found " + versions["sentence-transformers"]) + if "torch" in required: + import torch + versions["torch_runtime"] = torch.__version__ + print(json.dumps({"python": platform.python_version(), "platform": platform.platform(), "dependencies": versions})) + except Exception as error: + sys.exit("SentenceTransformers parity setup: " + str(error) + ". Install sentence-transformers==6.1.0, torch, numpy and datasets in a venv and set SENTENCE_TRANSFORMERS_PYTHON (or STING_PYTHON), or install uv.") + """ + + result = System.cmd(cmd, args ++ ["-c", code], stderr_to_stdout: true) + :persistent_term.put(key, result) + + if elem(result, 1) == 0 do + versions = + for app <- [:bumblebee, :axon, :nx], + into: %{}, + do: {app, to_string(Application.spec(app, :vsn))} + + IO.puts( + "SentenceTransformers parity runtime: " <> + Jason.encode!(%{ + elixir: System.version(), + otp: System.otp_release(), + backend: inspect(Nx.default_backend()), + dependencies: versions + }) + ) + + IO.puts("SentenceTransformers Python runtime: " <> String.trim(elem(result, 0))) + end + + result + + cached -> + cached + end + + case result do + {_, 0} -> :ok + {output, _} -> raise String.trim(output) + end + end + + defp python_command(deps) do + cond do + python_path = + System.get_env("SENTENCE_TRANSFORMERS_PYTHON") || System.get_env("STING_PYTHON") -> + {python_path, []} + + uv_path = System.find_executable("uv") -> + with_args = + Enum.flat_map(deps, fn dep -> + dep = if dep == "sentence-transformers", do: "sentence-transformers==6.1.0", else: dep + ["--with", dep] + end) + + {uv_path, ["run", "--quiet"] ++ with_args ++ ["python"]} + + py3_path = System.find_executable("python3") -> + {py3_path, []} + + true -> + raise "Neither 'uv' nor 'python3' executable was found in PATH" + end + end + + defp parse_result(output, opts) do + case Regex.run(~r/__RESULT__:(.*)$/m, output) do + [_, json_data] -> + data = Jason.decode!(json_data) + type = Keyword.get(opts, :type, :f32) + + if type == :raw do + data + else + Nx.tensor(data, type: type) + end + + nil -> + raise "Could not find __RESULT__ marker in Python output:\n#{output}" + end + end +end From b228b4ceed9b31d30d446d454e414d8d0360126d Mon Sep 17 00:00:00 2001 From: Daniel Kukula Date: Tue, 6 Oct 2026 22:33:25 +0200 Subject: [PATCH 2/4] Refactor SentenceTransformers modules and pipeline to reduce nesting --- lib/bumblebee/sentence_transformers.ex | 503 ++++++------ .../sentence_transformers/downloader.ex | 30 +- .../sentence_transformers/modules.ex | 678 ++++++++++++++++ .../sentence_transformers/pipeline.ex | 738 +++--------------- .../sentence_transformers/python_bridge.ex | 3 +- 5 files changed, 1023 insertions(+), 929 deletions(-) create mode 100644 lib/bumblebee/sentence_transformers/modules.ex diff --git a/lib/bumblebee/sentence_transformers.ex b/lib/bumblebee/sentence_transformers.ex index 431f8b9e..ca9ec1b1 100644 --- a/lib/bumblebee/sentence_transformers.ex +++ b/lib/bumblebee/sentence_transformers.ex @@ -54,76 +54,70 @@ defmodule Bumblebee.SentenceTransformers do |> Keyword.take([:backend, :type, :safetensors_reader]) |> Keyword.put_new(:architecture, :base) - case Downloader.download_file(repository, "modules.json") do - {:ok, modules_path} -> - with {:ok, modules} <- Pipeline.decode_json(modules_path), - modules = Enum.sort_by(modules, &(&1["idx"] || 0)), - [first | pipeline_modules] <- modules do - first_type = Pipeline.normalize_module_type(first["type"]) - - case first_type do - "Transformer" -> - base_repository = Downloader.subrepository(repository, first["path"]) - - with {:ok, model_info} <- Bumblebee.load_model(base_repository, base_opts), - {:ok, model_info} <- - Pipeline.build(repository, pipeline_modules, model_info, opts) do - load_config_and_attach(repository, model_info) - end - - "StaticEmbedding" -> - with {:ok, model_info} <- - Pipeline.build_static_embedding(repository, first, opts), - {:ok, model_info} <- - Pipeline.build(repository, pipeline_modules, model_info, opts) do - load_config_and_attach(repository, model_info) - end - - other -> - {:error, "unsupported base module in modules.json: #{inspect(other)}"} - end - else - [] -> - {:error, "modules.json is empty"} - - {:error, reason} -> - {:error, reason} - - _other -> - {:error, - "expected modules.json to define a Transformer or StaticEmbedding base model"} + case load_pipeline_modules(repository) do + {:ok, [first | pipeline_modules]} -> + with {:ok, model_info} <- load_base_model(repository, first, base_opts, opts), + {:ok, model_info} <- Pipeline.build(repository, pipeline_modules, model_info, opts) do + load_config_and_attach(repository, model_info) end + {:missing_modules, _reason} -> + load_fallback_model(repository, base_opts, opts) + {:error, reason} -> - if Downloader.missing_file?(repository, "modules.json", reason) do - load_fallback_model(repository, base_opts, opts) - else - {:error, reason} - end + {:error, reason} end end + defp load_base_model(repository, %{"path" => path, "type" => type} = first, base_opts, opts) do + case Pipeline.normalize_module_type(type) do + "Transformer" -> + base_repository = Downloader.subrepository(repository, path) + Bumblebee.load_model(base_repository, base_opts) + + "StaticEmbedding" -> + Pipeline.build_static_embedding(repository, first, opts) + + other -> + {:error, "unsupported base module in modules.json: #{inspect(other)}"} + end + end + + defp load_base_model(_repository, _first, _base_opts, _opts) do + {:error, "expected modules.json to define a Transformer or StaticEmbedding base model"} + end + defp load_fallback_model(repository, base_opts, opts) do # Fallback when modules.json is not present: load model and apply default pooling # CausalLM-based models use last token pooling, otherwise mean pooling - with {:ok, model_info} <- Bumblebee.load_model(repository, base_opts) do - pooling_mode = infer_fallback_pooling_mode(repository, model_info) - - with {:ok, model_info} <- - Pipeline.build( - repository, - [ - %{ - "type" => "Pooling", - "path" => "default_pooling", - "pooling_mode" => pooling_mode - } - ], - model_info, - opts - ) do - load_config_and_attach(repository, model_info) - end + with {:ok, model_info} <- Bumblebee.load_model(repository, base_opts), + fallback_modules = build_fallback_modules(repository, model_info), + {:ok, model_info} <- Pipeline.build(repository, fallback_modules, model_info, opts) do + load_config_and_attach(repository, model_info) + end + end + + defp build_fallback_modules(repository, model_info) do + pooling_mode = infer_fallback_pooling_mode(repository, model_info) + [%{"type" => "Pooling", "path" => "default_pooling", "pooling_mode" => pooling_mode}] + end + + defp load_pipeline_modules(repository, opts \\ []) do + case Downloader.download_file(repository, "modules.json", opts) do + {:ok, path} -> + with {:ok, modules} <- Pipeline.decode_json(path) do + case Enum.sort_by(modules, &(&1["idx"] || 0)) do + [] -> {:error, "modules.json is empty"} + sorted -> {:ok, sorted} + end + end + + {:error, reason} -> + if Downloader.missing_file?(repository, "modules.json", reason) do + {:missing_modules, reason} + else + {:error, reason} + end end end @@ -167,28 +161,22 @@ defmodule Bumblebee.SentenceTransformers do end defp do_load_embedding_head(repository, model_info, opts) do - case Downloader.download_file(repository, "modules.json") do - {:ok, modules_path} -> - with {:ok, modules} <- Pipeline.decode_json(modules_path), - modules = Enum.sort_by(modules, &(&1["idx"] || 0)), - [first | pipeline_modules] <- modules, - true <- - Pipeline.normalize_module_type(first["type"]) in ["Transformer", "StaticEmbedding"] || - {:error, - "expected the first module in modules.json to be Transformer or StaticEmbedding"}, - {:ok, model_info} <- Pipeline.build(repository, pipeline_modules, model_info, opts) do - load_config_and_attach(repository, model_info) + case load_pipeline_modules(repository) do + {:ok, [first | pipeline_modules]} -> + if Pipeline.normalize_module_type(first["type"]) in ["Transformer", "StaticEmbedding"] do + with {:ok, model_info} <- Pipeline.build(repository, pipeline_modules, model_info, opts) do + load_config_and_attach(repository, model_info) + end else - {:error, reason} -> - {:error, reason} - - _other -> - {:error, - "expected modules.json to define a Transformer or StaticEmbedding base model"} + {:error, + "expected the first module in modules.json to be Transformer or StaticEmbedding"} end - {:error, _} -> + {:missing_modules, _reason} -> {:error, "could not find modules.json in the repository"} + + {:error, reason} -> + {:error, reason} end end @@ -197,22 +185,15 @@ defmodule Bumblebee.SentenceTransformers do or the repository root when `modules.json` is absent. """ def load_tokenizer(repository, opts \\ []) do - case Downloader.download_file(repository, "modules.json") do - {:ok, path} -> - with {:ok, modules} <- Pipeline.decode_json(path), - [first | _] <- Enum.sort_by(modules, &(&1["idx"] || 0)) do - Bumblebee.load_tokenizer(Downloader.subrepository(repository, first["path"]), opts) - else - {:error, reason} -> {:error, reason} - _ -> {:error, "modules.json is empty"} - end + case load_pipeline_modules(repository, opts) do + {:ok, [first | _]} -> + Bumblebee.load_tokenizer(Downloader.subrepository(repository, first["path"]), opts) + + {:missing_modules, _reason} -> + Bumblebee.load_tokenizer(repository, opts) {:error, reason} -> - if Downloader.missing_file?(repository, "modules.json", reason) do - Bumblebee.load_tokenizer(repository, opts) - else - {:error, reason} - end + {:error, reason} end end @@ -424,88 +405,10 @@ defmodule Bumblebee.SentenceTransformers do {_init_fun, encoder} = Axon.build(model) embedding_fun = fn params, inputs -> - output = encoder.(params, inputs) - - output = - case output do - %Nx.Tensor{} -> - output - - %{^output_attribute => %Axon.None{}} -> - keys = output |> Map.keys() |> Enum.sort() - - raise ArgumentError, - "key #{inspect(output_attribute)} in the output map has value %Axon.None{}," <> - " make sure this is the correct key and check module documentation incase it is opt-in." <> - " The output map keys are: #{inspect(keys)}" - - %{^output_attribute => output} -> - output - - %{embedding: output} when output_attribute == :pooled_state -> - output - - %{} -> - keys = output |> Map.keys() |> Enum.sort() - - raise ArgumentError, - "key #{inspect(output_attribute)} not found in the output map," <> - " you may want to set :output_attribute to one of the map keys: #{inspect(keys)}" - - _ -> - output - end - - if output_pool != nil and Nx.rank(output) != 3 do - raise ArgumentError, - "expected the output tensor to have rank 3 to apply :output_pool, got: #{Nx.rank(output)}." <> - " You should either disable pooling or pick a different output using :output_attribute" - end - - output = - case output_pool do - nil -> - output - - :cls_token_pooling -> - Nx.take(output, 0, axis: 1) - - :mean_pooling -> - input_mask_expanded = Nx.new_axis(inputs["attention_mask"], -1) - - output - |> Nx.multiply(input_mask_expanded) - |> Nx.sum(axes: [1]) - |> Nx.divide(Nx.sum(input_mask_expanded, axes: [1])) - - :last_token_pooling -> - sequence_lengths = - inputs["attention_mask"] - |> Nx.sum(axes: [1]) - |> Nx.subtract(1) - |> Nx.as_type({:s, 64}) - - Bumblebee.Utils.Nx.batched_take(output, sequence_lengths) - - other -> - raise ArgumentError, - "expected :output_pool to be one of :cls_token_pooling, :mean_pooling, :last_token_pooling or nil, got: #{inspect(other)}" - end - - output = - case embedding_processor do - nil -> - output - - :l2_norm -> - Bumblebee.Utils.Nx.normalize(output) - - other -> - raise ArgumentError, - "expected :embedding_processor to be one of nil or :l2_norm, got: #{inspect(other)}" - end - - output + encoder.(params, inputs) + |> extract_output_tensor(output_attribute) + |> apply_output_pool(output_pool, inputs) + |> apply_embedding_processor(embedding_processor) end batch_keys = Bumblebee.Shared.sequence_batch_keys(sequence_length) @@ -622,36 +525,104 @@ defmodule Bumblebee.SentenceTransformers do |> MapSet.new(&PreTrainedTokenizer.token_to_id(tokenizer, &1)) end + defp extract_output_tensor(%Nx.Tensor{} = tensor, _attribute), do: tensor + + defp extract_output_tensor(%{} = output, attribute) do + case Map.fetch(output, attribute) do + {:ok, %Axon.None{}} -> + keys = output |> Map.keys() |> Enum.sort() + + raise ArgumentError, + "key #{inspect(attribute)} in the output map has value %Axon.None{}," <> + " make sure this is the correct key and check module documentation incase it is opt-in." <> + " The output map keys are: #{inspect(keys)}" + + {:ok, tensor} -> + tensor + + :error when attribute == :pooled_state and is_map_key(output, :embedding) -> + output.embedding + + :error -> + keys = output |> Map.keys() |> Enum.sort() + + raise ArgumentError, + "key #{inspect(attribute)} not found in the output map," <> + " you may want to set :output_attribute to one of the map keys: #{inspect(keys)}" + end + end + + defp extract_output_tensor(other, _attribute), do: other + + defp apply_output_pool(output, nil, _inputs), do: output + + defp apply_output_pool(output, pool, inputs) do + if Nx.rank(output) != 3 do + raise ArgumentError, + "expected the output tensor to have rank 3 to apply :output_pool, got: #{Nx.rank(output)}." <> + " You should either disable pooling or pick a different output using :output_attribute" + end + + case pool do + :cls_token_pooling -> + Nx.take(output, 0, axis: 1) + + :mean_pooling -> + input_mask_expanded = Nx.new_axis(inputs["attention_mask"], -1) + + output + |> Nx.multiply(input_mask_expanded) + |> Nx.sum(axes: [1]) + |> Nx.divide(Nx.sum(input_mask_expanded, axes: [1])) + + :last_token_pooling -> + sequence_lengths = + inputs["attention_mask"] + |> Nx.sum(axes: [1]) + |> Nx.subtract(1) + |> Nx.as_type({:s, 64}) + + Bumblebee.Utils.Nx.batched_take(output, sequence_lengths) + + other -> + raise ArgumentError, + "expected :output_pool to be one of :cls_token_pooling, :mean_pooling, :last_token_pooling or nil, got: #{inspect(other)}" + end + end + + defp apply_embedding_processor(output, nil), do: output + defp apply_embedding_processor(output, :l2_norm), do: Bumblebee.Utils.Nx.normalize(output) + + defp apply_embedding_processor(_output, other) do + raise ArgumentError, + "expected :embedding_processor to be one of nil or :l2_norm, got: #{inspect(other)}" + end + defp resolve_prompt(opts, model_info) do - case {opts[:prompt], opts[:prompt_name]} do - {prompt, nil} when is_binary(prompt) -> + case {Keyword.fetch(opts, :prompt), Keyword.fetch(opts, :prompt_name)} do + {{:ok, prompt}, :error} when is_binary(prompt) -> prompt - {prompt, nil} when prompt in [false, nil] -> - if Keyword.has_key?(opts, :prompt) do - "" - else - find_default_prompt(model_info) - end + {{:ok, prompt}, :error} when prompt in [false, nil] -> + "" - {prompt, nil} -> + {{:ok, prompt}, :error} -> raise ArgumentError, "expected :prompt to be a string or false, got: #{inspect(prompt)}" - {nil, prompt_name} when is_binary(prompt_name) -> + {:error, {:ok, prompt_name}} when is_binary(prompt_name) -> find_prompt_by_name!(model_info, prompt_name) - {nil, prompt_name} when prompt_name in [false, nil] -> - if Keyword.has_key?(opts, :prompt_name) do - "" - else - find_default_prompt(model_info) - end + {:error, {:ok, prompt_name}} when prompt_name in [false, nil] -> + "" - {nil, prompt_name} -> + {:error, {:ok, prompt_name}} -> raise ArgumentError, "expected :prompt_name to be a string or false, got: #{inspect(prompt_name)}" - {_prompt, _prompt_name} -> + {:error, :error} -> + find_default_prompt(model_info) + + {{:ok, _}, {:ok, _}} -> raise ArgumentError, "expected either :prompt or :prompt_name, but both were given" end end @@ -717,36 +688,27 @@ defmodule Bumblebee.SentenceTransformers do {:error, "could not find sentence transformers config"} else config = Enum.reduce(configs, %{}, &Map.merge(&2, &1)) - tokenizer_limit = (tokenizer || root_tokenizer || %{})["model_max_length"] - # Hugging Face uses enormous sentinel values for tokenizers without a limit. - tokenizer_limit = - if is_integer(tokenizer_limit) and tokenizer_limit > 0 and - tokenizer_limit < 1_000_000_000, - do: tokenizer_limit - - limit = - (transformer || %{})["max_seq_length"] || config["max_seq_length"] || tokenizer_limit - + limit = resolve_max_seq_length(transformer, config, tokenizer, root_tokenizer) {:ok, parse_st_config(Map.put(config, "max_seq_length", limit))} end end end - defp first_module_path(repository, opts) do - case Downloader.download_file(repository, "modules.json", opts) do - {:ok, path} -> - with {:ok, modules} <- Pipeline.decode_json(path), - [first | _] <- Enum.sort_by(modules, &(&1["idx"] || 0)) do - {:ok, first["path"] || ""} - else - {:error, reason} -> {:error, reason} - _ -> {:error, "modules.json is empty"} - end + defp resolve_max_seq_length(transformer, config, tokenizer, root_tokenizer) do + tokenizer_limit = (tokenizer || root_tokenizer || %{})["model_max_length"] - {:error, reason} -> - if Downloader.missing_file?(repository, "modules.json", reason), - do: {:ok, "0_Transformer"}, - else: {:error, reason} + tokenizer_limit = + if is_integer(tokenizer_limit) and tokenizer_limit > 0 and tokenizer_limit < 1_000_000_000, + do: tokenizer_limit + + (transformer || %{})["max_seq_length"] || config["max_seq_length"] || tokenizer_limit + end + + defp first_module_path(repository, opts) do + case load_pipeline_modules(repository, opts) do + {:ok, [first | _]} -> {:ok, first["path"] || ""} + {:missing_modules, _} -> {:ok, "0_Transformer"} + {:error, reason} -> {:error, reason} end end @@ -800,71 +762,68 @@ defmodule Bumblebee.SentenceTransformers do defp parse_similarity_fn_name(nil), do: nil defp parse_similarity_fn_name(other), do: other + @causal_spec_names [ + "Gemma", + "GemmaConfig", + "Gemma3Text", + "Gemma3TextConfig", + "Llama", + "LlamaConfig", + "Mistral", + "MistralConfig", + "Qwen3", + "Qwen3Config", + "SmolLm3", + "Smollm3Config", + "Phi", + "PhiConfig", + "Phi3", + "Phi3Config", + "Gpt2", + "Gpt2Config", + "GptBigCode", + "GptBigCodeConfig", + "GptNeoX", + "GptNeoXConfig" + ] + @doc """ Infers fallback pooling mode when `modules.json` is missing. Returns `"last"` for causal-LM architectures, and `"mean"` otherwise. """ def infer_fallback_pooling_mode(repository, model_info) do - is_causal_from_config = - case Downloader.download_file(repository, "config.json") do - {:ok, path} -> - case Pipeline.decode_json(path) do - {:ok, config} -> - archs = config["architectures"] || [] - - causal_arch? = - Enum.any?(archs, fn arch -> - is_binary(arch) and - (String.ends_with?(arch, "ForCausalLM") or - String.ends_with?(arch, "LMHeadModel")) - end) - - causal_arch? and Map.get(config, "is_causal", true) - - _ -> - false - end + if causal_from_config?(repository) or causal_from_spec?(model_info) do + "last" + else + "mean" + end + end - _ -> - false - end + defp causal_from_config?(repository) do + with {:ok, path} <- Downloader.download_file(repository, "config.json"), + {:ok, config} <- Pipeline.decode_json(path) do + archs = config["architectures"] || [] - is_causal_from_spec = - if spec = model_info[:spec] || model_info[:config] do - struct_name = spec.__struct__ |> Module.split() |> List.last() + causal_arch? = + Enum.any?(archs, fn arch -> + is_binary(arch) and + (String.ends_with?(arch, "ForCausalLM") or String.ends_with?(arch, "LMHeadModel")) + end) - struct_name in [ - "Gemma", - "GemmaConfig", - "Gemma3Text", - "Gemma3TextConfig", - "Llama", - "LlamaConfig", - "Mistral", - "MistralConfig", - "Qwen3", - "Qwen3Config", - "SmolLm3", - "Smollm3Config", - "Phi", - "PhiConfig", - "Phi3", - "Phi3Config", - "Gpt2", - "Gpt2Config", - "GptBigCode", - "GptBigCodeConfig", - "GptNeoX", - "GptNeoXConfig" - ] - else + causal_arch? and Map.get(config, "is_causal", true) + else + _ -> false + end + end + + defp causal_from_spec?(model_info) do + case model_info[:spec] || model_info[:config] do + nil -> false - end - if is_causal_from_config or is_causal_from_spec do - "last" - else - "mean" + spec -> + struct_name = spec.__struct__ |> Module.split() |> List.last() + struct_name in @causal_spec_names end end end diff --git a/lib/bumblebee/sentence_transformers/downloader.ex b/lib/bumblebee/sentence_transformers/downloader.ex index 27241fa5..a65b25b6 100644 --- a/lib/bumblebee/sentence_transformers/downloader.ex +++ b/lib/bumblebee/sentence_transformers/downloader.ex @@ -22,28 +22,12 @@ defmodule Bumblebee.SentenceTransformers.Downloader do def download_file(repository, filename, opts \\ []) def download_file({:local, dir}, filename, _opts) when is_binary(dir) do - path = Path.join(dir, filename) - - if File.exists?(path) do - {:ok, path} - else - {:error, "file #{inspect(path)} does not exist"} - end + check_local_file(Path.join(dir, filename)) end def download_file({:local, dir, local_opts}, filename, _opts) when is_binary(dir) do - path = - if subdir = local_opts[:subdir] do - Path.join([dir, subdir, filename]) - else - Path.join(dir, filename) - end - - if File.exists?(path) do - {:ok, path} - else - {:error, "file #{inspect(path)} does not exist"} - end + subdir = local_opts[:subdir] || "" + check_local_file(Path.join([dir, subdir, filename])) end def download_file({:hf, repo_id}, filename, opts) do @@ -74,6 +58,14 @@ defmodule Bumblebee.SentenceTransformers.Downloader do {:error, "invalid repository format: #{inspect(other)}"} end + defp check_local_file(path) do + if File.exists?(path) do + {:ok, path} + else + {:error, "file #{inspect(path)} does not exist"} + end + end + @doc false def missing_file?(repository, filename, reason) do case repository do diff --git a/lib/bumblebee/sentence_transformers/modules.ex b/lib/bumblebee/sentence_transformers/modules.ex new file mode 100644 index 00000000..e4c75fff --- /dev/null +++ b/lib/bumblebee/sentence_transformers/modules.ex @@ -0,0 +1,678 @@ +defmodule Bumblebee.SentenceTransformers.Modules do + @moduledoc false + + alias Axon.MixedPrecision.Policy + alias Bumblebee.SentenceTransformers.Downloader + alias Bumblebee.SentenceTransformers.Pipeline + + @weights_filenames ["model.safetensors", "pytorch_model.bin"] + + @doc """ + Builds a single pipeline module step on top of the current computation node. + """ + def build(repository, module, node, attention_mask, params_data, params_parameters, opts) do + case Pipeline.normalize_module_type(module["type"]) do + "Pooling" -> + build_pooling(repository, module, node, attention_mask, params_data, params_parameters) + + "Dense" -> + build_dense(repository, module, node, params_data, params_parameters, opts) + + "FusedDense" -> + build_fused_dense(repository, module, node, params_data, params_parameters, opts) + + "Normalize" -> + build_normalize(module, node, params_data, params_parameters) + + "Dropout" -> + build_dropout(node, params_data, params_parameters) + + "LayerNorm" -> + build_layer_norm(repository, module, node, params_data, params_parameters, opts) + + "WeightedLayerPooling" -> + build_weighted_layer_pooling( + repository, + module, + node, + params_data, + params_parameters, + opts + ) + + other -> + {:error, "unsupported SentenceTransformers module #{inspect(other)}"} + end + end + + @doc """ + Builds a Pooling module on top of the given node. + """ + def build_pooling( + repository, + module, + node, + attention_mask, + params_data, + params_parameters + ) do + with {:ok, config} <- load_pooling_config(repository, module), + {:ok, node} <- pooling_layer(node, attention_mask, config, module["path"]) do + requires_prompt_length = Map.get(config, "include_prompt", true) == false + {:ok, node, params_data, params_parameters, requires_prompt_length} + end + end + + defp load_pooling_config(_repository, %{"config" => config}) when is_map(config) do + {:ok, config} + end + + defp load_pooling_config(_repository, %{"path" => "default_pooling"} = module) do + mode = module["pooling_mode"] || "mean" + {:ok, %{"pooling_mode" => mode}} + end + + defp load_pooling_config(repository, %{"path" => path}) do + load_module_config(repository, path) + end + + @doc """ + Builds a Dense module on top of the given node. + """ + def build_dense( + repository, + %{"path" => path}, + node, + params_data, + params_parameters, + opts + ) do + with {:ok, config} <- load_module_config(repository, path), + {:ok, activation} <- parse_activation(config["activation_function"], path), + {:ok, weights_file, weights_path} <- find_weights_file(repository, path) do + in_features = config["in_features"] + out_features = config["out_features"] + bias? = Map.get(config, "bias", true) + use_residual? = Map.get(config, "use_residual", false) + + out_node = + node + |> Axon.dense(out_features, use_bias: bias?, name: path) + |> maybe_activation(activation) + + tensors = load_tensors(weights_file, weights_path, opts) + layer_params = load_dense_params(tensors, bias?, opts) + + {params_data, params_parameters} = + put_layer_params(params_data, params_parameters, path, layer_params) + + {node, params_data, params_parameters} = + maybe_dense_residual( + out_node, + node, + use_residual?, + in_features, + out_features, + tensors, + path, + params_data, + params_parameters, + opts + ) + + {:ok, node, params_data, params_parameters, false} + end + end + + defp load_dense_params(tensors, bias?, opts) do + kernel = + tensors["linear.weight"] + |> Nx.to_tensor() + |> Nx.transpose() + |> cast_param(opts[:type]) + |> allocate_param(opts[:backend]) + + 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 + end + + defp maybe_dense_residual( + out_node, + _node, + false, + _in_features, + _out_features, + _tensors, + _path, + params_data, + params_parameters, + _opts + ) do + {out_node, params_data, params_parameters} + end + + defp maybe_dense_residual( + out_node, + node, + true, + features, + features, + _tensors, + path, + params_data, + params_parameters, + _opts + ) do + {Axon.add(out_node, node, name: "#{path}.add_residual"), params_data, params_parameters} + end + + defp maybe_dense_residual( + out_node, + node, + true, + _in_features, + out_features, + tensors, + path, + params_data, + params_parameters, + opts + ) do + res_name = "#{path}.residual" + res_node = Axon.dense(node, out_features, use_bias: false, name: res_name) + + res_kernel = + tensors["residual.weight"] + |> Nx.to_tensor() + |> Nx.transpose() + |> cast_param(opts[:type]) + |> allocate_param(opts[:backend]) + + {params_data, params_parameters} = + put_layer_params(params_data, params_parameters, res_name, %{"kernel" => res_kernel}) + + {Axon.add(out_node, res_node, name: "#{path}.add_residual"), params_data, params_parameters} + end + + @doc """ + Builds a FusedDense module on top of the given node. + """ + def build_fused_dense( + repository, + module, + node, + params_data, + params_parameters, + opts + ) do + [p1, p2] = module["paths"] + layer_name = module["path"] + out_features = module["out_features"] + + with {:ok, w1_file, w1_path} <- find_weights_file(repository, p1), + {:ok, w2_file, w2_path} <- find_weights_file(repository, p2) 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 = + k1 + |> Nx.dot(k2) + |> cast_param(opts[:type]) + |> allocate_param(opts[:backend]) + + node = Axon.dense(node, out_features, use_bias: false, name: layer_name) + + {params_data, params_parameters} = + put_layer_params(params_data, params_parameters, layer_name, %{"kernel" => k_fused}) + + {:ok, node, params_data, params_parameters, false} + end + end + + @doc """ + Builds a Normalize module on top of the given node. + """ + def build_normalize(%{"path" => path}, node, params_data, params_parameters) do + node = Axon.nx(node, &Bumblebee.Utils.Nx.normalize/1, name: "#{path}.normalize") + {:ok, node, params_data, params_parameters, false} + end + + @doc """ + Builds a Dropout module (no-op at inference). + """ + def build_dropout(node, params_data, params_parameters) do + {:ok, node, params_data, params_parameters, false} + end + + @doc """ + Builds a LayerNorm module on top of the given node. + """ + def build_layer_norm( + repository, + %{"path" => path}, + node, + params_data, + params_parameters, + opts + ) do + with {:ok, config} <- load_module_config(repository, path), + {:ok, weights_file, weights_path} <- find_weights_file(repository, path) do + eps = config["eps"] || 1.0e-5 + layer_name = path + + node = Axon.layer_norm(node, epsilon: eps, name: layer_name) + tensors = load_tensors(weights_file, weights_path, opts) + + gamma = + (tensors["weight"] || tensors["gamma"] || tensors["norm.weight"]) + |> Nx.to_tensor() + |> cast_param(opts[:type]) + |> allocate_param(opts[:backend]) + + beta = + (tensors["bias"] || tensors["beta"] || tensors["norm.bias"]) + |> Nx.to_tensor() + |> cast_param(opts[:type]) + |> allocate_param(opts[:backend]) + + {params_data, params_parameters} = + put_layer_params(params_data, params_parameters, layer_name, %{ + "gamma" => gamma, + "beta" => beta + }) + + {:ok, node, params_data, params_parameters, false} + end + end + + @doc """ + Builds a WeightedLayerPooling module on top of the given node. + """ + def build_weighted_layer_pooling( + repository, + %{"path" => path}, + node, + params_data, + params_parameters, + opts + ) do + with {:ok, config} <- load_module_config(repository, path) do + layer_start = config["layer_start"] || 4 + num_hidden_layers = config["num_hidden_layers"] || 12 + num_weights = max(num_hidden_layers + 1 - layer_start, 1) + layer_name = path + + weights = load_layer_weights(repository, path, num_weights, opts) + + {params_data, params_parameters} = + put_layer_params(params_data, params_parameters, layer_name, %{"layer_weights" => weights}) + + param_layer = Axon.param("layer_weights", {num_weights}, initializer: :ones) + + weighted_node = + Axon.layer( + fn stacked_layers, layer_weights, _opts -> + total_layers = Nx.axis_size(stacked_layers, 0) + + selected = + Nx.slice_along_axis( + stacked_layers, + layer_start, + total_layers - layer_start, + axis: 0 + ) + + num_selected = total_layers - layer_start + weights_reshaped = Nx.reshape(layer_weights, {num_selected, 1, 1, 1}) + weighted = Nx.multiply(weights_reshaped, selected) + sum = Nx.sum(weighted, axes: [0]) + total_weight = Nx.sum(layer_weights) + Nx.divide(sum, total_weight) + end, + [node, param_layer], + name: layer_name + ) + + {:ok, weighted_node, params_data, params_parameters, false} + end + end + + defp load_layer_weights(repository, path, num_weights, opts) do + raw_weights = + case find_weights_file(repository, path) do + {:ok, weights_file, weights_path} -> + tensors = load_tensors(weights_file, weights_path, opts) + tensors["layer_weights"] || tensors["weight"] || Nx.broadcast(1.0, {num_weights}) + + _ -> + Nx.broadcast(1.0, {num_weights}) + end + + raw_weights + |> Nx.to_tensor() + |> cast_param(opts[:type]) + |> allocate_param(opts[:backend]) + end + + @doc """ + Builds a base model from a StaticEmbedding module (e.g. Model2Vec). + """ + def build_static_embedding(repository, %{"path" => path}, opts) do + with {:ok, weights_file, weights_path} <- find_weights_file(repository, path) do + tensors = load_tensors(weights_file, weights_path, opts) + + embedding_weight = + (tensors["embedding.weight"] || tensors["weight"]) + |> Nx.to_tensor() + |> cast_param(opts[:type]) + |> allocate_param(opts[:backend]) + + {vocab_size, embedding_dim} = Nx.shape(embedding_weight) + layer_name = "#{path}.embedding" + + model = + "input_ids" + |> Axon.input(optional: false) + |> Axon.embedding(vocab_size, embedding_dim, name: layer_name) + |> Axon.nx(fn hidden_state -> + %{hidden_state: hidden_state} + end) + + params = %Axon.ModelState{ + data: %{layer_name => %{"kernel" => embedding_weight}}, + parameters: %{layer_name => ["kernel"]}, + frozen_parameters: %{}, + state: %{} + } + + {:ok, %{model: model, params: params, spec: nil}} + end + end + + @doc """ + Constructs a default mean-pooling layer. + """ + def default_pooling(node, attention_mask, name \\ "default_pooling") do + config = %{"pooling_mode_mean_tokens" => true} + pooling_layer(node, attention_mask, config, name) + end + + @doc """ + Constructs a pooling layer from a configuration map. + """ + def pooling_layer(hidden_state, attention_mask, config, name) do + include_prompt = Map.get(config, "include_prompt", true) + + case extract_pooling_modes(config) do + [] -> + {:error, "no supported pooling mode found in #{name}"} + + modes -> + layer = build_pooling_layer(hidden_state, attention_mask, modes, include_prompt, name) + {:ok, layer} + end + end + + defp extract_pooling_modes(config) do + modes_from_flags = + 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_weightedmean_tokens"], :weightedmean_tokens}, + {["pooling_mode_lasttoken", "pooling_mode_last_token"], :last_token} + ], + Enum.any?(keys, &(config[&1] == true)), + do: mode + + case {modes_from_flags, config["pooling_mode"]} do + {[], "mean"} -> [:mean_tokens] + {[], "cls"} -> [:cls_token] + {[], "max"} -> [:max_tokens] + {[], "mean_sqrt_len_tokens"} -> [:mean_sqrt_len_tokens] + {[], "weightedmean"} -> [:weightedmean_tokens] + {[], "last"} -> [:last_token] + {[], "lasttoken"} -> [:last_token] + {modes, _} -> modes + end + end + + defp build_pooling_layer(hidden_state, attention_mask, modes, true, name) do + Axon.layer( + fn hidden_state, attention_mask, _opts -> + apply_pooling_modes(hidden_state, attention_mask, modes) + end, + [hidden_state, attention_mask], + name: name + ) + end + + defp build_pooling_layer(hidden_state, attention_mask, modes, false, name) do + prompt_length = Axon.input("prompt_length", optional: true) + + Axon.layer( + fn hidden_state, attention_mask, prompt_length, _opts -> + attention_mask = + case prompt_length do + %Axon.None{} -> attention_mask + _ -> exclude_prompt_from_mask(attention_mask, prompt_length) + end + + apply_pooling_modes(hidden_state, attention_mask, modes) + end, + [hidden_state, attention_mask, prompt_length], + name: name + ) + end + + defp apply_pooling_modes(hidden_state, attention_mask, modes) do + 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 = mask |> Nx.sum(axes: [1]) |> Nx.max(eps) + Nx.divide(sum_embeddings, sum_mask) + + :cls_token -> + first_indices = Nx.argmax(attention_mask, axis: 1) + Bumblebee.Utils.Nx.batched_take(hidden_state, first_indices) + + :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.neg_infinity(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 = mask |> Nx.sum(axes: [1]) |> Nx.max(eps) + Nx.divide(sum_embeddings, Nx.sqrt(sum_mask)) + + :weightedmean_tokens -> + mask = attention_mask |> Nx.as_type(type) |> Nx.new_axis(-1) + seq_len = Nx.axis_size(hidden_state, 1) + + weights = + {1, seq_len, 1} + |> Nx.iota(type: type) + |> Nx.add(1) + + weighted_mask = Nx.multiply(mask, weights) + sum_embeddings = Nx.sum(Nx.multiply(hidden_state, weighted_mask), axes: [1]) + sum_mask = weighted_mask |> Nx.sum(axes: [1]) |> Nx.max(eps) + Nx.divide(sum_embeddings, sum_mask) + + :last_token -> + lengths = + attention_mask + |> Nx.not_equal(0) + |> Nx.select( + Nx.iota(Nx.shape(attention_mask), axis: 1), + 0 + ) + |> Nx.reduce_max(axes: [1]) + |> Nx.as_type({:s, 64}) + + selected = Bumblebee.Utils.Nx.batched_take(hidden_state, lengths) + valid = attention_mask |> Nx.not_equal(0) |> Nx.any(axes: [1]) |> Nx.new_axis(-1) + Nx.select(Nx.broadcast(valid, Nx.shape(selected)), selected, 0) + end) + + case pool_outputs do + [single] -> single + multiple -> Nx.concatenate(multiple, axis: -1) + end + end + + defp exclude_prompt_from_mask(attention_mask, prompt_length) do + type = Nx.type(attention_mask) + seq_len = Nx.axis_size(attention_mask, 1) + batch_size = Nx.axis_size(attention_mask, 0) + + prompt_length = + case Nx.rank(prompt_length) do + 0 -> Nx.broadcast(Nx.as_type(prompt_length, type), {batch_size, 1}) + 1 -> prompt_length |> Nx.as_type(type) |> Nx.reshape({batch_size, 1}) + 2 -> Nx.as_type(prompt_length, type) + end + + positions = Nx.iota({1, seq_len}, type: type) + + pad_lengths = + attention_mask + |> Nx.as_type({:s, 64}) + |> Nx.argmax(axis: 1) + |> Nx.as_type(type) + |> Nx.reshape({batch_size, 1}) + + has_left_pad = Nx.any(Nx.not_equal(pad_lengths, 0)) + + threshold_right = prompt_length + threshold_left = Nx.add(pad_lengths, prompt_length) + + threshold = + Nx.select(Nx.broadcast(has_left_pad, {batch_size, 1}), threshold_left, threshold_right) + + keep_mask = Nx.greater_equal(positions, threshold) + + Nx.multiply(attention_mask, Nx.as_type(keep_mask, type)) + 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("torch.nn.modules.activation.Sigmoid", _path), do: {:ok, :sigmoid} + defp parse_activation("torch.nn.modules.activation.Mish", _path), do: {:ok, :mish} + defp parse_activation("torch.nn.modules.activation.LeakyReLU", _path), do: {:ok, :leaky_relu} + defp parse_activation("transformers.activations.GELUActivation", _path), do: {:ok, :gelu} + defp parse_activation("transformers.activations.FastGELUActivation", _path), do: {:ok, :gelu} + defp parse_activation("transformers.activations.NewGELUActivation", _path), do: {:ok, :gelu} + defp parse_activation("transformers.activations.SiLUActivation", _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) + + @doc false + def load_module_config(repository, path) do + with {:ok, config_path} <- + Downloader.download_file(repository, Path.join(path, "config.json")) do + decode_json(config_path) + end + end + + def find_weights_file(repository, dir) do + Enum.find_value(@weights_filenames, fn filename -> + path = Path.join(dir, filename) + + case Downloader.download_file(repository, path) do + {:ok, downloaded_path} -> {:ok, path, downloaded_path} + _ -> nil + end + end) || {:error, "could not find parameters file in #{dir}"} + end + + def 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 + + def cast_param(tensor, nil), do: tensor + + def cast_param(tensor, %Policy{params: type}) do + Nx.as_type(tensor, type) + end + + def cast_param(tensor, type) do + type = Nx.Type.normalize!(type) + Nx.as_type(tensor, type) + end + + def allocate_param(tensor, nil), do: tensor + + def allocate_param(tensor, backend) do + Nx.with_default_backend(backend, fn -> Nx.backend_copy(tensor) end) + end + + def 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 + + defp put_layer_params(params_data, params_parameters, layer_name, layer_params) do + 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 + + {params_data, params_parameters} + end +end diff --git a/lib/bumblebee/sentence_transformers/pipeline.ex b/lib/bumblebee/sentence_transformers/pipeline.ex index c4b5434b..b8ab1342 100644 --- a/lib/bumblebee/sentence_transformers/pipeline.ex +++ b/lib/bumblebee/sentence_transformers/pipeline.ex @@ -3,7 +3,15 @@ defmodule Bumblebee.SentenceTransformers.Pipeline do alias Axon.MixedPrecision.Policy alias Bumblebee.Layers - alias Bumblebee.SentenceTransformers.Downloader + alias Bumblebee.SentenceTransformers.Modules + + defdelegate pooling_layer(hidden_state, attention_mask, config, name), to: Modules + defdelegate default_pooling(node, attention_mask, name \\ "default_pooling"), to: Modules + defdelegate build_static_embedding(repository, module, opts), to: Modules + defdelegate decode_json(path), to: Modules + defdelegate cast_param(tensor, type), to: Modules + defdelegate allocate_param(tensor, backend), to: Modules + defdelegate load_module_config(repository, path), to: Modules @doc """ Builds the pipeline modules on top of base model_info. @@ -22,23 +30,7 @@ defmodule Bumblebee.SentenceTransformers.Pipeline do normalize_module_type(m["type"]) == "WeightedLayerPooling" end) - extracted_hidden_states = - if has_weighted_pooling?, do: extract_all_hidden_states(base_model) - - hidden_state = - if extracted_hidden_states do - extracted_hidden_states - else - Axon.nx(base_model, fn - %{hidden_state: hidden_state} -> - hidden_state - - output -> - raise ArgumentError, - "expected base model output to contain :hidden_state, but got keys: #{inspect(Map.keys(output))}." <> - " Ensure you loaded the base encoder model." - end) - end + hidden_state = extract_base_hidden_state(base_model, has_weighted_pooling?) attention_mask = Layers.default Axon.input("attention_mask", optional: true) do @@ -48,59 +40,9 @@ defmodule Bumblebee.SentenceTransformers.Pipeline do initial_state = {hidden_state, model_info.params.data, model_info.params.parameters, false} - result = - Enum.reduce_while( - modules, - {:ok, initial_state}, - fn module, {:ok, {node, params_data, params_parameters, prompt_length_required}} -> - case build_module( - repository, - module, - node, - attention_mask, - params_data, - params_parameters, - opts - ) do - {:ok, node, params_data, params_parameters, requires_prompt_length} -> - {:cont, - {:ok, - {node, params_data, params_parameters, - prompt_length_required or requires_prompt_length}}} - - {:error, reason} -> - {:halt, {:error, reason}} - end - end - ) - - with {:ok, {embedding, params_data, params_parameters, prompt_length_required}} <- result do - embedding = - case opts[:truncate_dim] do - truncate_dim when is_integer(truncate_dim) and truncate_dim > 0 -> - truncated = - Axon.nx( - embedding, - fn tensor -> Nx.slice_along_axis(tensor, 0, truncate_dim, axis: -1) end, - name: "truncate_dim" - ) - - needs_renorm? = - opts[:normalize] == true or - opts[:normalize_embeddings] == true or - pipeline_has_normalize?(modules) - - if needs_renorm? do - Axon.nx(truncated, &Bumblebee.Utils.Nx.normalize/1, - name: "truncate_dim.normalize" - ) - else - truncated - end - - _ -> - embedding - end + with {:ok, {embedding, params_data, params_parameters, prompt_length_required}} <- + build_pipeline_modules(repository, modules, initial_state, attention_mask, opts) do + embedding = maybe_truncate_embedding(embedding, modules, opts) final_model = Axon.container(%{ @@ -108,8 +50,7 @@ defmodule Bumblebee.SentenceTransformers.Pipeline do pooled_state: embedding, hidden_state: hidden_state }) - - final_model = apply_type(final_model, opts[:type]) + |> apply_type(opts[:type]) final_params = if params_parameters do @@ -141,461 +82,78 @@ defmodule Bumblebee.SentenceTransformers.Pipeline do def prompt_length_required?(%{prompt_length_required: value}), do: value def prompt_length_required?(_), do: false - defp build_module( - repository, - %{"path" => path} = module, - node, - attention_mask, - params_data, - params_parameters, - opts - ) do - case normalize_module_type(module["type"]) do - "Pooling" -> - config_result = - cond do - config = module["config"] -> - {:ok, config} - - path == "default_pooling" -> - mode = module["pooling_mode"] || "mean" - {:ok, %{"pooling_mode" => mode}} - - true -> - with {:ok, config_path} <- - Downloader.download_file(repository, Path.join(path, "config.json")) do - decode_json(config_path) - end - end - - with {:ok, config} <- config_result, - {:ok, node} <- pooling_layer(node, attention_mask, config, path) do - requires_prompt_length = Map.get(config, "include_prompt", true) == false - {:ok, node, params_data, params_parameters, requires_prompt_length} - end - - "Dense" -> - with {:ok, config_path} <- - Downloader.download_file(repository, Path.join(path, "config.json")), - {:ok, config} <- decode_json(config_path), - {:ok, activation} <- parse_activation(config["activation_function"], path), - {:ok, weights_file, weights_path} <- find_weights_file(repository, path) do - in_features = config["in_features"] - out_features = config["out_features"] - bias = Map.get(config, "bias", true) - use_residual = Map.get(config, "use_residual", false) - layer_name = path - - out_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 - - {node, params_data, params_parameters} = - if use_residual do - if in_features == out_features do - {Axon.add(out_node, node, name: "#{layer_name}.add_residual"), params_data, - params_parameters} - else - res_name = "#{layer_name}.residual" - res_node = Axon.dense(node, out_features, use_bias: false, name: res_name) - - res_kernel = - tensors["residual.weight"] - |> Nx.to_tensor() - |> Nx.transpose() - |> cast_param(opts[:type]) - |> allocate_param(opts[:backend]) - - res_params = %{"kernel" => res_kernel} - params_data = Map.put(params_data, res_name, res_params) - - params_parameters = - if params_parameters do - Map.put(params_parameters, res_name, ["kernel"]) - end - - {Axon.add(out_node, res_node, name: "#{layer_name}.add_residual"), params_data, - params_parameters} - end - else - {out_node, params_data, params_parameters} - end - - {:ok, node, params_data, params_parameters, false} - end - - "FusedDense" -> - [p1, p2] = module["paths"] - layer_name = module["path"] - out_features = module["out_features"] - - with {:ok, w1_file, w1_path} <- find_weights_file(repository, p1), - {:ok, w2_file, w2_path} <- find_weights_file(repository, p2) 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 = - k1 - |> Nx.dot(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, false} - end - - "Normalize" -> - node = Axon.nx(node, &Bumblebee.Utils.Nx.normalize/1, name: "#{path}.normalize") - {:ok, node, params_data, params_parameters, false} - - "Dropout" -> - {:ok, node, params_data, params_parameters, false} - - "LayerNorm" -> - with {:ok, config_path} <- - Downloader.download_file(repository, Path.join(path, "config.json")), - {:ok, config} <- decode_json(config_path), - {:ok, weights_file, weights_path} <- find_weights_file(repository, path) do - eps = config["eps"] || 1.0e-5 - layer_name = path - - node = Axon.layer_norm(node, epsilon: eps, name: layer_name) - tensors = load_tensors(weights_file, weights_path, opts) - - gamma = - (tensors["weight"] || tensors["gamma"] || tensors["norm.weight"]) - |> Nx.to_tensor() - |> cast_param(opts[:type]) - |> allocate_param(opts[:backend]) - - beta = - (tensors["bias"] || tensors["beta"] || tensors["norm.bias"]) - |> Nx.to_tensor() - |> cast_param(opts[:type]) - |> allocate_param(opts[:backend]) - - layer_params = %{"gamma" => gamma, "beta" => beta} - params_data = Map.put(params_data, layer_name, layer_params) - - params_parameters = - if params_parameters do - Map.put(params_parameters, layer_name, ["gamma", "beta"]) - end - - {:ok, node, params_data, params_parameters, false} - end - - "WeightedLayerPooling" -> - with {:ok, config_path} <- - Downloader.download_file(repository, Path.join(path, "config.json")), - {:ok, config} <- decode_json(config_path) do - layer_start = config["layer_start"] || 4 - num_hidden_layers = config["num_hidden_layers"] || 12 - num_weights = max(num_hidden_layers + 1 - layer_start, 1) - layer_name = path - - {_weights, params_data, params_parameters} = - case find_weights_file(repository, path) do - {:ok, weights_file, weights_path} -> - tensors = load_tensors(weights_file, weights_path, opts) - - w = - (tensors["layer_weights"] || tensors["weight"] || - Nx.broadcast(1.0, {num_weights})) - |> Nx.to_tensor() - |> cast_param(opts[:type]) - |> allocate_param(opts[:backend]) - - params_data = Map.put(params_data, layer_name, %{"layer_weights" => w}) - - params_parameters = - if params_parameters do - Map.put(params_parameters, layer_name, ["layer_weights"]) - end - - {w, params_data, params_parameters} - - _ -> - w = - 1.0 - |> Nx.broadcast({num_weights}) - |> cast_param(opts[:type]) - |> allocate_param(opts[:backend]) - - params_data = Map.put(params_data, layer_name, %{"layer_weights" => w}) - - params_parameters = - if params_parameters do - Map.put(params_parameters, layer_name, ["layer_weights"]) - end - - {w, params_data, params_parameters} - end - - param_layer = Axon.param("layer_weights", {num_weights}, initializer: :ones) - - weighted_node = - Axon.layer( - fn stacked_layers, layer_weights, _opts -> - total_layers = Nx.axis_size(stacked_layers, 0) - - selected = - Nx.slice_along_axis( - stacked_layers, - layer_start, - total_layers - layer_start, - axis: 0 - ) - - num_selected = total_layers - layer_start - weights_reshaped = Nx.reshape(layer_weights, {num_selected, 1, 1, 1}) - weighted = Nx.multiply(weights_reshaped, selected) - sum = Nx.sum(weighted, axes: [0]) - total_weight = Nx.sum(layer_weights) - Nx.divide(sum, total_weight) - end, - [node, param_layer], - name: layer_name - ) - - {:ok, weighted_node, params_data, params_parameters, false} - end - - other -> - {:error, "unsupported SentenceTransformers module #{inspect(other)}"} - end + defp extract_base_hidden_state(base_model, true) do + extract_all_hidden_states(base_model) || extract_single_hidden_state(base_model) end - def default_pooling(node, attention_mask, name \\ "default_pooling") do - config = %{"pooling_mode_mean_tokens" => true} - pooling_layer(node, attention_mask, config, name) + defp extract_base_hidden_state(base_model, false) do + extract_single_hidden_state(base_model) end - def pooling_layer(hidden_state, attention_mask, config, name) do - include_prompt = Map.get(config, "include_prompt", true) - - modes_from_flags = - 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_weightedmean_tokens"], :weightedmean_tokens}, - {["pooling_mode_lasttoken", "pooling_mode_last_token"], :last_token} - ], - Enum.any?(keys, &(config[&1] == true)), - do: mode - - modes = - case {modes_from_flags, config["pooling_mode"]} do - {[], "mean"} -> [:mean_tokens] - {[], "cls"} -> [:cls_token] - {[], "max"} -> [:max_tokens] - {[], "mean_sqrt_len_tokens"} -> [:mean_sqrt_len_tokens] - {[], "weightedmean"} -> [:weightedmean_tokens] - {[], "last"} -> [:last_token] - {[], "lasttoken"} -> [:last_token] - {modes, _} -> modes - end + defp extract_single_hidden_state(base_model) do + Axon.nx(base_model, fn + %{hidden_state: hidden_state} -> + hidden_state - if modes == [] do - {:error, "no supported pooling mode found in #{name}"} - else - layer = build_pooling_layer(hidden_state, attention_mask, modes, include_prompt, name) - {:ok, layer} - end - end - - defp build_pooling_layer(hidden_state, attention_mask, modes, true, name) do - Axon.layer( - fn hidden_state, attention_mask, _opts -> - apply_pooling_modes(hidden_state, attention_mask, modes) - end, - [hidden_state, attention_mask], - name: name - ) + output -> + raise ArgumentError, + "expected base model output to contain :hidden_state, but got keys: #{inspect(Map.keys(output))}." <> + " Ensure you loaded the base encoder model." + end) end - defp build_pooling_layer(hidden_state, attention_mask, modes, false, name) do - prompt_length = Axon.input("prompt_length", optional: true) - - Axon.layer( - fn hidden_state, attention_mask, prompt_length, _opts -> - attention_mask = - case prompt_length do - %Axon.None{} -> attention_mask - _ -> exclude_prompt_from_mask(attention_mask, prompt_length) - end - - apply_pooling_modes(hidden_state, attention_mask, modes) - end, - [hidden_state, attention_mask, prompt_length], - name: name + defp build_pipeline_modules(repository, modules, initial_state, attention_mask, opts) do + Enum.reduce_while( + modules, + {:ok, initial_state}, + fn module, {:ok, {node, params_data, params_parameters, prompt_length_required}} -> + case Modules.build( + repository, + module, + node, + attention_mask, + params_data, + params_parameters, + opts + ) do + {:ok, node, params_data, params_parameters, requires_prompt_length} -> + {:cont, + {:ok, + {node, params_data, params_parameters, + prompt_length_required or requires_prompt_length}}} + + {:error, reason} -> + {:halt, {:error, reason}} + end + end ) end - defp apply_pooling_modes(hidden_state, attention_mask, modes) do - 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 = mask |> Nx.sum(axes: [1]) |> Nx.max(eps) - Nx.divide(sum_embeddings, sum_mask) - - :cls_token -> - first_indices = Nx.argmax(attention_mask, axis: 1) - Bumblebee.Utils.Nx.batched_take(hidden_state, first_indices) - - :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.neg_infinity(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 = mask |> Nx.sum(axes: [1]) |> Nx.max(eps) - Nx.divide(sum_embeddings, Nx.sqrt(sum_mask)) - - :weightedmean_tokens -> - mask = attention_mask |> Nx.as_type(type) |> Nx.new_axis(-1) - seq_len = Nx.axis_size(hidden_state, 1) - - weights = - {1, seq_len, 1} - |> Nx.iota(type: type) - |> Nx.add(1) - - weighted_mask = Nx.multiply(mask, weights) - sum_embeddings = Nx.sum(Nx.multiply(hidden_state, weighted_mask), axes: [1]) - sum_mask = weighted_mask |> Nx.sum(axes: [1]) |> Nx.max(eps) - Nx.divide(sum_embeddings, sum_mask) - - :last_token -> - lengths = - attention_mask - |> Nx.not_equal(0) - |> Nx.select( - Nx.iota(Nx.shape(attention_mask), axis: 1), - 0 - ) - |> Nx.reduce_max(axes: [1]) - |> Nx.as_type({:s, 64}) - - selected = Bumblebee.Utils.Nx.batched_take(hidden_state, lengths) - valid = attention_mask |> Nx.not_equal(0) |> Nx.any(axes: [1]) |> Nx.new_axis(-1) - Nx.select(Nx.broadcast(valid, Nx.shape(selected)), selected, 0) - end) + defp maybe_truncate_embedding(embedding, modules, opts) do + case opts[:truncate_dim] do + truncate_dim when is_integer(truncate_dim) and truncate_dim > 0 -> + truncated = + Axon.nx( + embedding, + fn tensor -> Nx.slice_along_axis(tensor, 0, truncate_dim, axis: -1) end, + name: "truncate_dim" + ) + + if needs_renorm?(modules, opts) do + Axon.nx(truncated, &Bumblebee.Utils.Nx.normalize/1, name: "truncate_dim.normalize") + else + truncated + end - case pool_outputs do - [single] -> single - multiple -> Nx.concatenate(multiple, axis: -1) + _ -> + embedding end end - defp exclude_prompt_from_mask(attention_mask, prompt_length) do - type = Nx.type(attention_mask) - seq_len = Nx.axis_size(attention_mask, 1) - batch_size = Nx.axis_size(attention_mask, 0) - - prompt_length = - case Nx.rank(prompt_length) do - 0 -> Nx.broadcast(Nx.as_type(prompt_length, type), {batch_size, 1}) - 1 -> prompt_length |> Nx.as_type(type) |> Nx.reshape({batch_size, 1}) - 2 -> Nx.as_type(prompt_length, type) - end - - positions = Nx.iota({1, seq_len}, type: type) - - pad_lengths = - attention_mask - |> Nx.as_type({:s, 64}) - |> Nx.argmax(axis: 1) - |> Nx.as_type(type) - |> Nx.reshape({batch_size, 1}) - - has_left_pad = Nx.any(Nx.not_equal(pad_lengths, 0)) - - threshold_right = prompt_length - threshold_left = Nx.add(pad_lengths, prompt_length) - - threshold = - Nx.select(Nx.broadcast(has_left_pad, {batch_size, 1}), threshold_left, threshold_right) - - keep_mask = Nx.greater_equal(positions, threshold) - - Nx.multiply(attention_mask, Nx.as_type(keep_mask, type)) - 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("torch.nn.modules.activation.Sigmoid", _path), do: {:ok, :sigmoid} - defp parse_activation("torch.nn.modules.activation.Mish", _path), do: {:ok, :mish} - defp parse_activation("torch.nn.modules.activation.LeakyReLU", _path), do: {:ok, :leaky_relu} - defp parse_activation("transformers.activations.GELUActivation", _path), do: {:ok, :gelu} - defp parse_activation("transformers.activations.FastGELUActivation", _path), do: {:ok, :gelu} - defp parse_activation("transformers.activations.NewGELUActivation", _path), do: {:ok, :gelu} - defp parse_activation("transformers.activations.SiLUActivation", _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}"} + defp needs_renorm?(modules, opts) do + opts[:normalize] == true or + opts[:normalize_embeddings] == true or + pipeline_has_normalize?(modules) end defp pipeline_has_normalize?(modules) do @@ -604,56 +162,53 @@ defmodule Bumblebee.SentenceTransformers.Pipeline do end) end - defp maybe_activation(node, nil), do: node - defp maybe_activation(node, activation), do: Axon.activation(node, activation) - defp maybe_fuse_dense(repository, modules, true) do fuse_dense_modules(repository, modules, []) end defp maybe_fuse_dense(_repository, modules, _false), do: {:ok, modules} - defp fuse_dense_modules( - repository, - [%{"type" => t1, "path" => p1} = m1, %{"type" => t2, "path" => p2} = m2 | rest], - acc - ) do - if normalize_module_type(t1) == "Dense" and normalize_module_type(t2) == "Dense" do - with {:ok, c1} <- load_module_config(repository, p1), - {:ok, c2} <- load_module_config(repository, p2), - 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(repository, rest, [fused_module | acc]) - else - _ -> - fuse_dense_modules(repository, [m2 | rest], [m1 | acc]) - end - else - fuse_dense_modules(repository, [m2 | rest], [m1 | acc]) + defp fuse_dense_modules(repository, [m1, m2 | rest], acc) do + case try_fuse_dense(repository, m1, m2) do + {:ok, fused} -> fuse_dense_modules(repository, rest, [fused | acc]) + :error -> fuse_dense_modules(repository, [m2 | rest], [m1 | acc]) end end - defp fuse_dense_modules(repository, [other | rest], acc) do - fuse_dense_modules(repository, rest, [other | acc]) + defp fuse_dense_modules(_repository, [other], acc) do + {:ok, Enum.reverse([other | acc])} end defp fuse_dense_modules(_repository, [], acc) do {:ok, Enum.reverse(acc)} end - defp load_module_config(repository, path) do - with {:ok, config_path} <- - Downloader.download_file(repository, Path.join(path, "config.json")) do - decode_json(config_path) + defp try_fuse_dense( + repository, + %{"type" => t1, "path" => p1}, + %{"type" => t2, "path" => p2} = _m2 + ) do + if normalize_module_type(t1) == "Dense" and normalize_module_type(t2) == "Dense" do + with {:ok, c1} <- Modules.load_module_config(repository, p1), + {:ok, c2} <- Modules.load_module_config(repository, p2), + true <- can_fuse_dense?(c1, c2) do + {:ok, + %{ + "type" => "sentence_transformers.models.FusedDense", + "path" => "#{p1}_#{p2}", + "paths" => [p1, p2], + "out_features" => c2["out_features"] + }} + else + _ -> :error + end + else + :error end end + defp try_fuse_dense(_repository, _m1, _m2), do: :error + 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"] @@ -663,49 +218,6 @@ defmodule Bumblebee.SentenceTransformers.Pipeline do c1_linear? and c2_linear? and no_bias? and dims_match? end - @weights_filenames ["model.safetensors", "pytorch_model.bin"] - - defp find_weights_file(repository, dir) do - Enum.find_value(@weights_filenames, fn filename -> - path = Path.join(dir, filename) - - case Downloader.download_file(repository, path) do - {:ok, downloaded_path} -> {:ok, path, downloaded_path} - _ -> nil - end - end) || {:error, "could not find parameters file in #{dir}"} - 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 - - @doc false - def cast_param(tensor, nil), do: tensor - - def cast_param(tensor, %Policy{params: type}) do - Nx.as_type(tensor, type) - end - - def cast_param(tensor, type) do - type = Nx.Type.normalize!(type) - Nx.as_type(tensor, type) - end - - @doc false - def allocate_param(tensor, nil), do: tensor - - def 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, %Policy{} = policy) do @@ -742,54 +254,6 @@ defmodule Bumblebee.SentenceTransformers.Pipeline do defp find_first_tensor(_, _fun), do: nil - def 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 - - @doc """ - Builds a base model from a StaticEmbedding module (e.g. Model2Vec). - """ - def build_static_embedding(repository, %{"path" => path}, opts) do - with {:ok, weights_file, weights_path} <- find_weights_file(repository, path) do - tensors = load_tensors(weights_file, weights_path, opts) - - embedding_weight = - (tensors["embedding.weight"] || tensors["weight"]) - |> Nx.to_tensor() - |> cast_param(opts[:type]) - |> allocate_param(opts[:backend]) - - {vocab_size, embedding_dim} = Nx.shape(embedding_weight) - layer_name = "#{path}.embedding" - - model = - "input_ids" - |> Axon.input(optional: false) - |> Axon.embedding(vocab_size, embedding_dim, name: layer_name) - |> Axon.nx(fn hidden_state -> - %{hidden_state: hidden_state} - end) - - params = %Axon.ModelState{ - data: %{layer_name => %{"kernel" => embedding_weight}}, - parameters: %{layer_name => ["kernel"]}, - frozen_parameters: %{}, - state: %{} - } - - {:ok, %{model: model, params: params, spec: nil}} - end - end - defp extract_all_hidden_states(%Axon{nodes: nodes}) do opt_in_node = Enum.find_value(nodes, fn {_id, node} -> diff --git a/test/support/sentence_transformers/python_bridge.ex b/test/support/sentence_transformers/python_bridge.ex index aea025db..6b16eb21 100644 --- a/test/support/sentence_transformers/python_bridge.ex +++ b/test/support/sentence_transformers/python_bridge.ex @@ -418,7 +418,8 @@ defmodule Bumblebee.SentenceTransformers.TestSupport.PythonBridge do [_, result_json] = Regex.run(~r/__RESULT__:(.*)$/m, output) inputs = - Regex.run(~r/__INPUTS__:(.*)$/m, output) + ~r/__INPUTS__:(.*)$/m + |> Regex.run(output) |> Enum.at(1) |> Jason.decode!() |> Map.new(fn {k, v} -> {k, Nx.tensor(v)} end) From 9edb175fdd32b7cea1516951fc7a4e573cbc176c Mon Sep 17 00:00:00 2001 From: Daniel Kukula Date: Tue, 6 Oct 2026 22:39:40 +0200 Subject: [PATCH 3/4] Remove Python bridge and runtime comparison tests --- .../python_comparison_test.exs | 557 ----------------- .../python_environment_test.exs | 48 -- .../sentence_transformers/python_bridge.ex | 558 ------------------ 3 files changed, 1163 deletions(-) delete mode 100644 test/bumblebee/sentence_transformers/python_comparison_test.exs delete mode 100644 test/bumblebee/sentence_transformers/python_environment_test.exs delete mode 100644 test/support/sentence_transformers/python_bridge.ex diff --git a/test/bumblebee/sentence_transformers/python_comparison_test.exs b/test/bumblebee/sentence_transformers/python_comparison_test.exs deleted file mode 100644 index f98e103b..00000000 --- a/test/bumblebee/sentence_transformers/python_comparison_test.exs +++ /dev/null @@ -1,557 +0,0 @@ -defmodule Bumblebee.SentenceTransformers.PythonComparisonTest do - use ExUnit.Case, async: false - - import Bumblebee.TestHelpers - - alias Bumblebee.SentenceTransformers - alias Bumblebee.SentenceTransformers.Pipeline - alias Bumblebee.SentenceTransformers.TestSupport.PythonBridge - - @moduletag :slow - @moduletag :python - - test "saved short token limit matches Python with root config, prompts and compilation" do - text = Enum.join(List.duplicate("hello world", 12), " ") - - for prompt <- [nil, "Represent the query: "] do - assert {:ok, %{dir: dir, python_output: expected}} = - PythonBridge.generate_random_sentence_transformer( - max_seq_length: 8, - text: text, - prompt: prompt - ) - - try do - assert {:ok, info} = SentenceTransformers.load_model({:local, dir}) - assert info.sentence_transformers.max_seq_length == 8 - assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) - - for opts <- [[], [compile: [batch_size: 1, sequence_length: 32]]] do - actual = - Nx.Serving.run(SentenceTransformers.text_embedding(info, tokenizer, opts), text).embedding - - assert_all_close(actual, Nx.squeeze(expected), atol: 1.0e-4) - end - after - File.rm_rf!(dir) - end - end - end - - test "pooling matches Python with left padding and fully masked sequences" do - hidden = {3, 5, 4} |> Nx.iota(type: :f32) |> Nx.add(1) - mask = Nx.tensor([[0, 0, 1, 1, 1], [1, 1, 0, 0, 0], [0, 0, 0, 0, 0]]) - modes = ["mean", "cls", "max", "mean_sqrt_len_tokens", "weightedmean", "lasttoken"] - - code = """ - import json - import torch - from sentence_transformers.models import Pooling - hidden = torch.tensor(#{Jason.encode!(Nx.to_list(hidden))}) - mask = torch.tensor(#{Jason.encode!(Nx.to_list(mask))}) - outputs = {} - for mode in #{inspect(modes)}: - result = Pooling(4, pooling_mode=mode)({"token_embeddings": hidden.clone(), "attention_mask": mask.clone()})["sentence_embedding"] - outputs[mode] = {"values": torch.nan_to_num(result, neginf=0).tolist(), "negative_infinity": torch.isneginf(result).int().tolist()} - print("__RESULT__:" + json.dumps(outputs)) - """ - - expected = PythonBridge.run_python_code(code, ["sentence-transformers", "torch"], type: :raw) - - for mode <- modes do - assert {:ok, model} = - Pipeline.pooling_layer( - Axon.input("hidden"), - Axon.input("mask"), - %{"pooling_mode" => mode}, - "pool" - ) - - {init, predict} = Axon.build(model) - inputs = %{"hidden" => hidden, "mask" => mask} - params = init.(inputs, Axon.ModelState.empty()) - actual = predict.(params, inputs) - infinite = Nx.is_infinity(actual) - assert Nx.to_list(infinite) == expected[mode]["negative_infinity"] - - assert_all_close(Nx.select(infinite, 0, actual), Nx.tensor(expected[mode]["values"]), - atol: 1.0e-4 - ) - end - end - - test "runtime prompts override saved defaults and work with compiled prompt exclusion" do - assert {:ok, %{dir: dir}} = - PythonBridge.generate_random_sentence_transformer( - prompt: "Represent the query: ", - include_prompt: false - ) - - try do - assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) - assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) - - for {prompt, compile} <- [ - {"Represent the document: ", nil}, - {"Represent the document: ", [batch_size: 2, sequence_length: 32]}, - {"", nil} - ] do - opts = [prompt: prompt] - opts = if compile, do: Keyword.put(opts, :compile, compile), else: opts - - actual = - model_info - |> SentenceTransformers.text_embedding(tokenizer, opts) - |> Nx.Serving.run("Hello world") - - expected = PythonBridge.run_sentence_transformers(dir, "Hello world", prompt: prompt) - assert_all_close(actual.embedding, expected, atol: 1.0e-4) - end - after - File.rm_rf!(dir) - end - end - - test "query and corpus helpers fall back when no named prompts are saved" do - assert {:ok, %{dir: dir, python_output: expected}} = - PythonBridge.generate_random_sentence_transformer() - - try do - assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) - assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) - - for helper <- [:encode_queries, :encode_corpus] do - [actual] = apply(SentenceTransformers, helper, [model_info, tokenizer, ["Hello world"]]) - assert_all_close(actual.embedding, expected, atol: 1.0e-4) - end - after - File.rm_rf!(dir) - end - end - - describe "On-the-fly generated random SentenceTransformers comparison" do - @pooling_modes [ - "mean_tokens", - "cls_token", - "max_tokens", - "mean_sqrt_len_tokens", - "weightedmean_tokens", - "last_token" - ] - - for mode <- @pooling_modes do - @mode mode - - test "matches Python for random model with #{@mode} pooling" do - assert {:ok, %{dir: dir, python_output: python_output, text: text}} = - PythonBridge.generate_random_sentence_transformer( - pooling_mode: @mode, - with_normalize: false - ) - - try do - assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) - assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) - - serving = SentenceTransformers.text_embedding(model_info, tokenizer) - res = Nx.Serving.run(serving, text) - - assert Nx.shape(res.embedding) == Nx.shape(python_output) - assert_all_close(res.embedding, python_output, atol: 1.0e-4) - after - File.rm_rf!(dir) - end - end - end - - test "matches Python for random model with mean pooling and normalization" do - assert {:ok, %{dir: dir, python_output: python_output, text: text}} = - PythonBridge.generate_random_sentence_transformer( - pooling_mode: "mean_tokens", - with_normalize: true - ) - - try do - assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) - assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) - - serving = SentenceTransformers.text_embedding(model_info, tokenizer) - res = Nx.Serving.run(serving, text) - - assert Nx.shape(res.embedding) == Nx.shape(python_output) - assert_all_close(res.embedding, python_output, atol: 1.0e-4) - after - File.rm_rf!(dir) - end - end - - test "matches Python for random model with Dense projection layer" do - assert {:ok, %{dir: dir, python_output: python_output, text: text}} = - PythonBridge.generate_random_sentence_transformer( - pooling_mode: "mean_tokens", - with_dense: true, - with_normalize: true - ) - - try do - assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) - assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) - - serving = SentenceTransformers.text_embedding(model_info, tokenizer) - res = Nx.Serving.run(serving, text) - - assert Nx.shape(res.embedding) == Nx.shape(python_output) - assert_all_close(res.embedding, python_output, atol: 1.0e-4) - after - File.rm_rf!(dir) - end - end - - test "matches Python for random model with prompt and include_prompt=false" do - prompt = "Represent the sentence for retrieval: " - - assert {:ok, %{dir: dir, python_output: python_output, text: text}} = - PythonBridge.generate_random_sentence_transformer( - pooling_mode: "mean_tokens", - prompt: prompt, - include_prompt: false, - text: "Hello world" - ) - - try do - assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) - assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) - - serving = SentenceTransformers.text_embedding(model_info, tokenizer) - res = Nx.Serving.run(serving, text) - - assert Nx.shape(res.embedding) == Nx.shape(python_output) - assert_all_close(res.embedding, python_output, atol: 1.0e-4) - after - File.rm_rf!(dir) - end - end - - test "matches Python for random model with prompt and include_prompt=true" do - prompt = "Represent the sentence for retrieval: " - - assert {:ok, %{dir: dir, python_output: python_output, text: text}} = - PythonBridge.generate_random_sentence_transformer( - pooling_mode: "mean_tokens", - prompt: prompt, - include_prompt: true, - text: "Hello world" - ) - - try do - assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) - assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) - - serving = SentenceTransformers.text_embedding(model_info, tokenizer) - res = Nx.Serving.run(serving, text) - - assert Nx.shape(res.embedding) == Nx.shape(python_output) - assert_all_close(res.embedding, python_output, atol: 1.0e-4) - after - File.rm_rf!(dir) - end - end - - test "matches Python encode_queries using config_sentence_transformers.json prompt" do - prompt = "Represent the query for retrieval: " - text = "Hello world" - - assert {:ok, %{dir: dir, python_output: python_output}} = - PythonBridge.generate_random_sentence_transformer( - pooling_mode: "mean_tokens", - prompt: prompt, - prompt_name: "query", - text: text - ) - - try do - assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) - assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) - - [res] = SentenceTransformers.encode_queries(model_info, tokenizer, [text]) - - assert Nx.shape(res.embedding) == Nx.shape(python_output) - assert_all_close(res.embedding, python_output, atol: 1.0e-4) - after - File.rm_rf!(dir) - end - end - - test "matches Python encode_corpus using config_sentence_transformers.json prompt" do - query_prompt = "Represent the query for retrieval: " - document_prompt = "Represent the document for retrieval: " - text = "Hello world" - - assert {:ok, %{dir: dir, python_output: python_output}} = - PythonBridge.generate_random_sentence_transformer( - pooling_mode: "mean_tokens", - prompt: query_prompt, - document_prompt: document_prompt, - prompt_name: "document", - text: text - ) - - try do - assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) - assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) - - [res] = SentenceTransformers.encode_corpus(model_info, tokenizer, [text]) - - assert Nx.shape(res.embedding) == Nx.shape(python_output) - assert_all_close(res.embedding, python_output, atol: 1.0e-4) - after - File.rm_rf!(dir) - end - end - end - - describe "On-the-fly generated Dense / Normalize / LayerNorm / Dropout comparison" do - test "matches Python for Dense with Tanh activation" do - text = "Hello world" - - assert {:ok, %{dir: dir, python_output: python_output}} = - PythonBridge.generate_random_sentence_transformer( - pooling_mode: "mean_tokens", - with_dense: true, - dense_activation: "Tanh", - text: text - ) - - try do - assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) - assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) - - serving = SentenceTransformers.text_embedding(model_info, tokenizer) - res = Nx.Serving.run(serving, text) - - assert Nx.shape(res.embedding) == Nx.shape(python_output) - assert_all_close(res.embedding, python_output, atol: 1.0e-4) - after - File.rm_rf!(dir) - end - end - - for activation <- ["Identity", "ReLU", "GELU", "SiLU", "Sigmoid", "Mish", "LeakyReLU"] do - @activation activation - - test "matches Python for Dense with #{@activation} activation" do - text = "Hello world" - - assert {:ok, %{dir: dir, python_output: python_output}} = - PythonBridge.generate_random_sentence_transformer( - pooling_mode: "mean_tokens", - with_dense: true, - dense_activation: @activation, - text: text - ) - - try do - assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) - assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) - - serving = SentenceTransformers.text_embedding(model_info, tokenizer) - res = Nx.Serving.run(serving, text) - - assert Nx.shape(res.embedding) == Nx.shape(python_output) - assert_all_close(res.embedding, python_output, atol: 1.0e-4) - after - File.rm_rf!(dir) - end - end - end - - test "matches Python for Normalize module" do - text = "Hello world" - - assert {:ok, %{dir: dir, python_output: python_output}} = - PythonBridge.generate_random_sentence_transformer( - pooling_mode: "mean_tokens", - with_normalize: true, - text: text - ) - - try do - assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) - assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) - - serving = SentenceTransformers.text_embedding(model_info, tokenizer) - res = Nx.Serving.run(serving, text) - - assert Nx.shape(res.embedding) == Nx.shape(python_output) - assert_all_close(res.embedding, python_output, atol: 1.0e-4) - after - File.rm_rf!(dir) - end - end - - test "matches Python for LayerNorm module" do - text = "Hello world" - - assert {:ok, %{dir: dir, python_output: python_output}} = - PythonBridge.generate_random_sentence_transformer( - pooling_mode: "mean_tokens", - with_layernorm: true, - text: text - ) - - try do - assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) - assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) - - serving = SentenceTransformers.text_embedding(model_info, tokenizer) - res = Nx.Serving.run(serving, text) - - assert Nx.shape(res.embedding) == Nx.shape(python_output) - assert_all_close(res.embedding, python_output, atol: 1.0e-4) - after - File.rm_rf!(dir) - end - end - - test "matches Python for Dropout module in eval mode" do - text = "Hello world" - - assert {:ok, %{dir: dir, python_output: python_output}} = - PythonBridge.generate_random_sentence_transformer( - pooling_mode: "mean_tokens", - with_dropout: true, - text: text - ) - - try do - assert {:ok, model_info} = SentenceTransformers.load_model({:local, dir}) - assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) - - serving = SentenceTransformers.text_embedding(model_info, tokenizer) - res = Nx.Serving.run(serving, text) - - assert Nx.shape(res.embedding) == Nx.shape(python_output) - assert_all_close(res.embedding, python_output, atol: 1.0e-4) - after - File.rm_rf!(dir) - end - end - - test "matches Python when fusing consecutive linear Dense layers" do - text = "Hello world" - - assert {:ok, %{dir: dir, python_output: python_output}} = - PythonBridge.generate_random_sentence_transformer( - pooling_mode: "mean_tokens", - with_fused_dense: true, - text: text - ) - - try do - assert {:ok, model_info} = - SentenceTransformers.load_model({:local, dir}, fuse_dense: true) - - assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) - - serving = SentenceTransformers.text_embedding(model_info, tokenizer) - res = Nx.Serving.run(serving, text) - - assert Nx.shape(res.embedding) == Nx.shape(python_output) - assert_all_close(res.embedding, python_output, atol: 1.0e-4) - after - File.rm_rf!(dir) - end - end - - for {description, opts} <- [ - {"first layer has bias", [fused_dense_first_bias: true]}, - {"first layer has activation", [fused_dense_first_activation: "Tanh"]}, - {"second layer has bias", [fused_dense_second_bias: true]}, - {"second layer has activation", [fused_dense_second_activation: "Tanh"]} - ] do - @description description - @opts opts - - test "refuses to fuse when #{@description}" do - text = "Hello world" - - assert {:ok, %{dir: dir, python_output: python_output}} = - PythonBridge.generate_random_sentence_transformer( - [pooling_mode: "mean_tokens", with_fused_dense: true, text: text] ++ @opts - ) - - try do - assert {:ok, model_info} = - SentenceTransformers.load_model({:local, dir}, fuse_dense: true) - - assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:local, dir}) - - serving = SentenceTransformers.text_embedding(model_info, tokenizer) - res = Nx.Serving.run(serving, text) - - assert Nx.shape(res.embedding) == Nx.shape(python_output) - assert_all_close(res.embedding, python_output, atol: 1.0e-4) - after - File.rm_rf!(dir) - end - end - end - end - - describe "HuggingFace Hub SentenceTransformers comparison" do - @tag :network - test "matches Python for sentence-transformers/all-MiniLM-L6-v2" do - repo_id = "sentence-transformers/all-MiniLM-L6-v2" - input_text = "Hello world" - - # 1. Elixir SentenceTransformers execution - assert {:ok, model_info} = SentenceTransformers.load_model({:hf, repo_id}) - assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer({:hf, repo_id}) - - serving = SentenceTransformers.text_embedding(model_info, tokenizer) - elixir_res = Nx.Serving.run(serving, input_text) - - # 2. Python sentence_transformers execution - python_res = PythonBridge.run_sentence_transformers(repo_id, input_text) - - # 3. Assertions and comparison - assert Nx.shape(elixir_res.embedding) == Nx.shape(python_res) - assert_all_close(elixir_res.embedding, python_res, atol: 1.0e-4) - end - end - - describe "HuggingFace tiny-random base models comparison" do - @tiny_random_models [ - {"hf-internal-testing/tiny-random-BertModel", "last_hidden_state"}, - {"hf-internal-testing/tiny-random-RobertaModel", "last_hidden_state"} - ] - - for {repo_id, output_key} <- @tiny_random_models do - @repo_id repo_id - @output_key output_key - - @tag :network - test "matches Python transformers for #{@repo_id}" do - inputs = %{ - "input_ids" => [[10, 20, 30, 40, 50]], - "attention_mask" => [[1, 1, 1, 1, 1]] - } - - elixir_inputs = %{ - "input_ids" => Nx.tensor(inputs["input_ids"]), - "attention_mask" => Nx.tensor(inputs["attention_mask"]) - } - - assert {:ok, %{model: model, params: params}} = Bumblebee.load_model({:hf, @repo_id}) - elixir_outputs = Axon.predict(model, params, elixir_inputs) - - python_tensor = - PythonBridge.run_hf_model(@repo_id, inputs, output_key: @output_key) - - elixir_tensor = Map.fetch!(elixir_outputs, :hidden_state) - assert Nx.shape(elixir_tensor) == Nx.shape(python_tensor) - assert_all_close(elixir_tensor, python_tensor, atol: 1.0e-4) - end - end - end -end diff --git a/test/bumblebee/sentence_transformers/python_environment_test.exs b/test/bumblebee/sentence_transformers/python_environment_test.exs deleted file mode 100644 index 22262437..00000000 --- a/test/bumblebee/sentence_transformers/python_environment_test.exs +++ /dev/null @@ -1,48 +0,0 @@ -defmodule Bumblebee.SentenceTransformers.PythonEnvironmentTest do - use ExUnit.Case, async: false - - alias Bumblebee.SentenceTransformers.TestSupport.PythonBridge - - @tag :python - @tag :tmp_dir - test "incompatible Python is checked once and reports an actionable version error", %{ - tmp_dir: dir - } do - python = System.find_executable("python3") - assert python, "python3 is needed for the Python environment validation regression" - wrapper = Path.join(dir, "incompatible-python") - count_file = Path.join(dir, "invocations") - - File.write!(wrapper, """ - #!#{python} - import importlib.metadata, sys - with open(#{Jason.encode!(count_file)}, "a") as counter: - counter.write("checked\\n") - importlib.metadata.version = lambda package: "0.0.0" - exec(sys.argv[-1]) - """) - - File.chmod!(wrapper, 0o755) - previous = System.get_env("SENTENCE_TRANSFORMERS_PYTHON") || System.get_env("STING_PYTHON") - System.put_env("SENTENCE_TRANSFORMERS_PYTHON", wrapper) - - on_exit(fn -> - if previous, - do: System.put_env("SENTENCE_TRANSFORMERS_PYTHON", previous), - else: System.delete_env("SENTENCE_TRANSFORMERS_PYTHON") - end) - - for _ <- 1..2 do - assert_raise RuntimeError, - ~r/expected sentence-transformers==6.1.0, found 0.0.0.*(SENTENCE_TRANSFORMERS_PYTHON|STING_PYTHON)/, - fn -> - PythonBridge.run_python_code( - "raise RuntimeError('must not execute')", - ["sentence-transformers"] - ) - end - end - - assert File.read!(count_file) == "checked\n" - end -end diff --git a/test/support/sentence_transformers/python_bridge.ex b/test/support/sentence_transformers/python_bridge.ex deleted file mode 100644 index 6b16eb21..00000000 --- a/test/support/sentence_transformers/python_bridge.ex +++ /dev/null @@ -1,558 +0,0 @@ -defmodule Bumblebee.SentenceTransformers.TestSupport.PythonBridge do - @moduledoc """ - Helper module for executing Python inference scripts (via `uv` or `python3`) - and generating random SentenceTransformers / Hugging Face model checkpoints - to compare outputs with Bumblebee.SentenceTransformers. - """ - - @doc """ - Runs inference using Python's `sentence_transformers` library on a given model and text, - returning the computed embeddings as an `Nx.Tensor`. - """ - def run_sentence_transformers(repo_id, text, opts \\ []) do - code = """ - import json - import os - import sys - import warnings - warnings.filterwarnings("ignore") - - os.environ["TOKENIZERS_PARALLELISM"] = "false" - os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" - - from sentence_transformers import SentenceTransformer - - model = SentenceTransformer("#{repo_id}") - encode_opts = json.loads(#{inspect(Jason.encode!(Map.new(Keyword.take(opts, [:prompt, :prompt_name]))))}) - res = model.encode(#{inspect(text)}, **encode_opts) - print("__RESULT__:" + json.dumps(res.tolist())) - """ - - deps = ["sentence-transformers", "torch"] - run_python(code, deps, opts) - end - - @pooling_mode_flags %{ - "mean_tokens" => "pooling_mode_mean_tokens=True", - "cls_token" => "pooling_mode_cls_token=True", - "max_tokens" => "pooling_mode_max_tokens=True", - "mean_sqrt_len_tokens" => "pooling_mode_mean_sqrt_len_tokens=True", - "weightedmean_tokens" => "pooling_mode_weightedmean_tokens=True", - "last_token" => "pooling_mode_lasttoken=True" - } - - @doc """ - Generates an on-the-fly random SentenceTransformer model checkpoint in Python with random weights, - saves it to a temporary directory, executes inference on `text`, and returns: - `{:ok, %{dir: dir, python_output: nx_tensor, text: text}}`. - - ## Options - * `:pooling_mode` - a single pooling mode name or a list of mode names. - Supported: `"mean_tokens"`, `"cls_token"`, `"max_tokens"`, - `"mean_sqrt_len_tokens"`, `"weightedmean_tokens"`, `"last_token"` - (default: `["mean_tokens"]`). - * `:with_dense` - whether to include a Dense projection layer (default: `false`) - * `:dense_activation` - activation for the Dense layer: `"Identity"`, `"Tanh"`, - `"ReLU"`, `"GELU"`, `"SiLU"`, `"Sigmoid"`, `"Mish"`, `"LeakyReLU"` - (default: `"Tanh"`) - * `:with_normalize` - whether to append a Normalize module (default: `false`) - * `:with_layernorm` - whether to append a LayerNorm module (default: `false`) - * `:with_dropout` - whether to append a Dropout module (default: `false`) - * `:with_fused_dense` - whether to append two consecutive Dense layers - (default: `false`). By default both use `Identity` activation and no bias so - Bumblebee.SentenceTransformers can fuse them. - * `:fused_dense_first_bias` - whether the first fused Dense layer has bias - (default: `false`) - * `:fused_dense_first_activation` - activation for the first fused Dense layer - (default: `"Identity"`) - * `:fused_dense_second_bias` - whether the second fused Dense layer has bias - (default: `false`) - * `:fused_dense_second_activation` - activation for the second fused Dense layer - (default: `"Identity"`) - * `:prompt` - prompt string to prepend to the input (default: `nil`) - * `:prompt_name` - name of a prompt in `config_sentence_transformers.json` to - use for the Python encode call (default: `nil`) - * `:document_prompt` - prompt string to register under the `"document"` name - (default: `nil`) - * `:include_prompt` - whether to include prompt tokens in pooling - (default: `true`). Only relevant when `:prompt` is given. - * `:text` - text to encode (default: `"Hello world"`) - """ - def generate_random_sentence_transformer(opts \\ []) do - pooling_mode = Keyword.get(opts, :pooling_mode, ["mean_tokens"]) - with_dense = Keyword.get(opts, :with_dense, false) - dense_activation = Keyword.get(opts, :dense_activation, "Tanh") - with_normalize = Keyword.get(opts, :with_normalize, false) - with_layernorm = Keyword.get(opts, :with_layernorm, false) - with_dropout = Keyword.get(opts, :with_dropout, false) - with_fused_dense = Keyword.get(opts, :with_fused_dense, false) - fused_dense_first_bias = Keyword.get(opts, :fused_dense_first_bias, false) - fused_dense_first_activation = Keyword.get(opts, :fused_dense_first_activation, "Identity") - fused_dense_second_bias = Keyword.get(opts, :fused_dense_second_bias, false) - fused_dense_second_activation = Keyword.get(opts, :fused_dense_second_activation, "Identity") - prompt = Keyword.get(opts, :prompt) - include_prompt = Keyword.get(opts, :include_prompt, true) - text = Keyword.get(opts, :text, "Hello world") - max_seq_length = Keyword.get(opts, :max_seq_length, 64) - - mode_flags = - pooling_mode - |> List.wrap() - |> Enum.map_join(", ", &Map.fetch!(@pooling_mode_flags, &1)) - - mode_flags = - if "mean_tokens" in List.wrap(pooling_mode), - do: mode_flags, - else: "pooling_mode_mean_tokens=False, " <> mode_flags - - pooling_args = - if prompt do - "#{mode_flags}, include_prompt=#{if include_prompt, do: "True", else: "False"}" - else - mode_flags - end - - prompt_name = Keyword.get(opts, :prompt_name) - - encode_call = - cond do - prompt_name -> - "st_model.encode(#{inspect(text)}, prompt_name=#{inspect(prompt_name)})" - - prompt -> - "st_model.encode(#{inspect(text)}, prompt=#{inspect(prompt)})" - - true -> - "st_model.encode(#{inspect(text)})" - end - - prompts = - if prompt do - %{"query" => prompt} - else - %{} - end - - prompts = - if document_prompt = Keyword.get(opts, :document_prompt) do - Map.put(prompts, "document", document_prompt) - else - prompts - end - - st_config = - if map_size(prompts) > 0 do - Jason.encode!(%{ - "prompts" => prompts, - "default_prompt_name" => "query" - }) - else - Jason.encode!(%{}) - end - - code = """ - import json - import os - import sys - import tempfile - import warnings - warnings.filterwarnings("ignore") - - os.environ["TOKENIZERS_PARALLELISM"] = "false" - os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" - - import torch - from transformers import BertConfig, BertModel, BertTokenizerFast - from sentence_transformers import SentenceTransformer, models - - torch.manual_seed(42) - - hidden_size = 32 - config = BertConfig( - vocab_size=128, - hidden_size=hidden_size, - num_hidden_layers=2, - num_attention_heads=4, - intermediate_size=64, - max_position_embeddings=64 - ) - base_model = BertModel(config) - - with tempfile.TemporaryDirectory() as tmp_base: - base_model.save_pretrained(tmp_base) - vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]", "hello", "world", - "represent", "the", "sentence", "for", "retrieval", ":", "query", "document"] - vocab += [f"unused_{i}" for i in range(128 - len(vocab))] - with open(os.path.join(tmp_base, "vocab.txt"), "w") as f: - f.write("\\n".join(vocab)) - BertTokenizerFast(vocab={token: i for i, token in enumerate(vocab)}, do_lower_case=True).save_pretrained(tmp_base) - - word_embedding_model = models.Transformer(tmp_base, max_seq_length=#{max_seq_length}) - pooling_model = models.Pooling( - word_embedding_model.get_embedding_dimension(), - #{pooling_args} - ) - - modules = [word_embedding_model, pooling_model] - - activation_map = { - "Identity": torch.nn.Identity(), - "Tanh": torch.nn.Tanh(), - "ReLU": torch.nn.ReLU(), - "GELU": torch.nn.GELU(), - "SiLU": torch.nn.SiLU(), - "Sigmoid": torch.nn.Sigmoid(), - "Mish": torch.nn.Mish(), - "LeakyReLU": torch.nn.LeakyReLU() - } - - if #{if with_dense, do: "True", else: "False"}: - dense_model = models.Dense( - in_features=hidden_size, - out_features=16, - bias=True, - activation_function=activation_map[#{inspect(dense_activation)}] - ) - modules.append(dense_model) - - if #{if with_fused_dense, do: "True", else: "False"}: - modules.append(models.Dense( - in_features=hidden_size, - out_features=24, - bias=#{if fused_dense_first_bias, do: "True", else: "False"}, - activation_function=activation_map[#{inspect(fused_dense_first_activation)}] - )) - modules.append(models.Dense( - in_features=24, - out_features=16, - bias=#{if fused_dense_second_bias, do: "True", else: "False"}, - activation_function=activation_map[#{inspect(fused_dense_second_activation)}] - )) - - if #{if with_normalize, do: "True", else: "False"}: - modules.append(models.Normalize()) - - if #{if with_layernorm, do: "True", else: "False"}: - modules.append(models.LayerNorm(hidden_size)) - - if #{if with_dropout, do: "True", else: "False"}: - modules.append(models.Dropout(0.5)) - - st_model = SentenceTransformer(modules=modules) - - save_dir = tempfile.mkdtemp(prefix="st_random_st_") - st_model.save(save_dir) - - with open(save_dir + "/config_sentence_transformers.json", "w") as f: - f.write(#{inspect(st_config)}) - - st_model = SentenceTransformer(save_dir) - - emb = #{encode_call}.tolist() - - print("__DIR__:" + save_dir) - print("__RESULT__:" + json.dumps(emb)) - """ - - deps = ["sentence-transformers", "torch"] - {cmd, args} = python_exec_args(code, deps) - - case System.cmd(cmd, args, stderr_to_stdout: true) do - {output, 0} -> - [_, dir] = Regex.run(~r/__DIR__:(.*)$/m, output) - [_, result_json] = Regex.run(~r/__RESULT__:(.*)$/m, output) - - python_output = Nx.tensor(Jason.decode!(result_json), type: :f32) - {:ok, %{dir: String.trim(dir), python_output: python_output, text: text}} - - {output, exit_code} -> - raise "Failed to generate random SentenceTransformer (exit #{exit_code}):\n#{output}" - end - end - - @doc """ - Runs inference using Python's `transformers` library on a given model and inputs. - """ - def run_hf_model(repo_id, inputs, opts \\ []) do - output_key = Keyword.get(opts, :output_key, "last_hidden_state") - json_inputs = Jason.encode!(inputs) - - code = """ - import json - import os - import sys - import warnings - warnings.filterwarnings("ignore") - - os.environ["TOKENIZERS_PARALLELISM"] = "false" - os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" - - import torch - from transformers import AutoModel - - raw_inputs = json.loads(#{inspect(json_inputs)}) - inputs = {k: torch.tensor(v) for k, v in raw_inputs.items()} - - model = AutoModel.from_pretrained("#{repo_id}") - model.eval() - - with torch.no_grad(): - outputs = model(**inputs) - result = getattr(outputs, "#{output_key}", None) - if result is None and "#{output_key}" in outputs: - result = outputs["#{output_key}"] - if hasattr(result, "cpu"): - result = result.cpu().numpy().tolist() - - print("__RESULT__:" + json.dumps(result)) - """ - - deps = ["transformers", "torch"] - run_python(code, deps, opts) - end - - @doc """ - Generates a tiny model with completely random weights on the fly in Python, saves it to - a temporary directory, executes inference in Python, and returns: - `%{dir: temp_dir, inputs: inputs_map, python_output: nx_tensor}`. - """ - def generate_random_hf_model(model_family, opts \\ []) do - seq_len = Keyword.get(opts, :seq_len, 5) - batch_size = Keyword.get(opts, :batch_size, 1) - - code = """ - import json - import os - import sys - import tempfile - import warnings - warnings.filterwarnings("ignore") - - os.environ["TOKENIZERS_PARALLELISM"] = "false" - os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" - - import torch - from transformers import BertConfig, BertModel, RobertaConfig, RobertaModel, GPT2Config, GPT2Model - - family = "#{model_family}" - batch_size = #{batch_size} - seq_len = #{seq_len} - - torch.manual_seed(42) - - if family == "bert": - config = BertConfig( - vocab_size=128, - hidden_size=32, - num_hidden_layers=2, - num_attention_heads=4, - intermediate_size=64, - max_position_embeddings=64 - ) - model = BertModel(config) - inputs = { - "input_ids": torch.randint(0, 128, (batch_size, seq_len)), - "attention_mask": torch.ones((batch_size, seq_len), dtype=torch.long) - } - output_attr = "last_hidden_state" - - elif family == "roberta": - config = RobertaConfig( - vocab_size=128, - hidden_size=32, - num_hidden_layers=2, - num_attention_heads=4, - intermediate_size=64, - max_position_embeddings=64, - type_vocab_size=1 - ) - model = RobertaModel(config) - inputs = { - "input_ids": torch.randint(0, 128, (batch_size, seq_len)), - "attention_mask": torch.ones((batch_size, seq_len), dtype=torch.long) - } - output_attr = "last_hidden_state" - - elif family == "gpt2": - config = GPT2Config( - vocab_size=128, - n_embd=32, - n_layer=2, - n_head=4, - n_positions=64, - pad_token_id=0, - bos_token_id=1, - eos_token_id=2 - ) - model = GPT2Model(config) - inputs = { - "input_ids": torch.randint(0, 128, (batch_size, seq_len)), - "attention_mask": torch.ones((batch_size, seq_len), dtype=torch.long) - } - output_attr = "last_hidden_state" - else: - raise ValueError(f"Unsupported family: {family}") - - model.eval() - - with torch.no_grad(): - outputs = model(**inputs) - result = getattr(outputs, output_attr).cpu().numpy().tolist() - - temp_dir = tempfile.mkdtemp(prefix="st_random_hf_") - model.save_pretrained(temp_dir) - - raw_inputs = {k: v.tolist() for k, v in inputs.items()} - - print("__DIR__:" + temp_dir) - print("__INPUTS__:" + json.dumps(raw_inputs)) - print("__RESULT__:" + json.dumps(result)) - """ - - deps = ["transformers", "torch"] - {cmd, args} = python_exec_args(code, deps) - - case System.cmd(cmd, args, stderr_to_stdout: true) do - {output, 0} -> - [_, dir] = Regex.run(~r/__DIR__:(.*)$/m, output) - [_, result_json] = Regex.run(~r/__RESULT__:(.*)$/m, output) - - inputs = - ~r/__INPUTS__:(.*)$/m - |> Regex.run(output) - |> Enum.at(1) - |> Jason.decode!() - |> Map.new(fn {k, v} -> {k, Nx.tensor(v)} end) - - python_output = Nx.tensor(Jason.decode!(result_json), type: :f32) - - {:ok, %{dir: String.trim(dir), inputs: inputs, python_output: python_output}} - - {output, exit_code} -> - raise "Failed to generate random model (exit #{exit_code}):\n#{output}" - end - end - - @doc """ - Runs custom Python code that prints `__RESULT__:` and returns the parsed Nx tensor or raw data. - """ - def run_python_code(code, deps \\ ["sentence-transformers", "torch"], opts \\ []) do - {cmd, args} = python_exec_args(code, deps) - - case System.cmd(cmd, args, stderr_to_stdout: true) do - {output, 0} -> - parse_result(output, opts) - - {output, exit_code} -> - raise "Python process failed with exit code #{exit_code}:\n#{output}" - end - end - - defp run_python(code, deps, opts) do - run_python_code(code, deps, opts) - end - - defp python_exec_args(code, deps) do - {cmd, args} = python_command(deps) - verify_environment!(cmd, args, deps) - {cmd, args ++ ["-c", code]} - end - - defp verify_environment!(cmd, args, deps) do - key = {__MODULE__, :environment, cmd, args, deps} - - result = - case :persistent_term.get(key, :unchecked) do - :unchecked -> - code = """ - import importlib.metadata as metadata - import json, platform, sys - required = #{Jason.encode!(deps)} - versions = {} - try: - for package in required: - versions[package] = metadata.version(package) - if "sentence-transformers" in versions and versions["sentence-transformers"] != "6.1.0": - raise RuntimeError("expected sentence-transformers==6.1.0, found " + versions["sentence-transformers"]) - if "torch" in required: - import torch - versions["torch_runtime"] = torch.__version__ - print(json.dumps({"python": platform.python_version(), "platform": platform.platform(), "dependencies": versions})) - except Exception as error: - sys.exit("SentenceTransformers parity setup: " + str(error) + ". Install sentence-transformers==6.1.0, torch, numpy and datasets in a venv and set SENTENCE_TRANSFORMERS_PYTHON (or STING_PYTHON), or install uv.") - """ - - result = System.cmd(cmd, args ++ ["-c", code], stderr_to_stdout: true) - :persistent_term.put(key, result) - - if elem(result, 1) == 0 do - versions = - for app <- [:bumblebee, :axon, :nx], - into: %{}, - do: {app, to_string(Application.spec(app, :vsn))} - - IO.puts( - "SentenceTransformers parity runtime: " <> - Jason.encode!(%{ - elixir: System.version(), - otp: System.otp_release(), - backend: inspect(Nx.default_backend()), - dependencies: versions - }) - ) - - IO.puts("SentenceTransformers Python runtime: " <> String.trim(elem(result, 0))) - end - - result - - cached -> - cached - end - - case result do - {_, 0} -> :ok - {output, _} -> raise String.trim(output) - end - end - - defp python_command(deps) do - cond do - python_path = - System.get_env("SENTENCE_TRANSFORMERS_PYTHON") || System.get_env("STING_PYTHON") -> - {python_path, []} - - uv_path = System.find_executable("uv") -> - with_args = - Enum.flat_map(deps, fn dep -> - dep = if dep == "sentence-transformers", do: "sentence-transformers==6.1.0", else: dep - ["--with", dep] - end) - - {uv_path, ["run", "--quiet"] ++ with_args ++ ["python"]} - - py3_path = System.find_executable("python3") -> - {py3_path, []} - - true -> - raise "Neither 'uv' nor 'python3' executable was found in PATH" - end - end - - defp parse_result(output, opts) do - case Regex.run(~r/__RESULT__:(.*)$/m, output) do - [_, json_data] -> - data = Jason.decode!(json_data) - type = Keyword.get(opts, :type, :f32) - - if type == :raw do - data - else - Nx.tensor(data, type: type) - end - - nil -> - raise "Could not find __RESULT__ marker in Python output:\n#{output}" - end - end -end From f39726cb37fedb37d4c8ba6a875dbf9e090cf271 Mon Sep 17 00:00:00 2001 From: Daniel Kukula Date: Tue, 6 Oct 2026 23:00:44 +0200 Subject: [PATCH 4/4] Add end-to-end integration test with intfloat/e5-small-v2 --- test/bumblebee/sentence_transformers_test.exs | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/test/bumblebee/sentence_transformers_test.exs b/test/bumblebee/sentence_transformers_test.exs index 4c2a209f..de8f7714 100644 --- a/test/bumblebee/sentence_transformers_test.exs +++ b/test/bumblebee/sentence_transformers_test.exs @@ -1286,4 +1286,30 @@ defmodule Bumblebee.SentenceTransformersTest do assert_in_delta val, ref, 1.0e-5 end) end + + @tag :slow + @tag :network + test "end-to-end integration test with intfloat/e5-small-v2" do + repo = {:hf, "intfloat/e5-small-v2"} + + assert {:ok, model_info} = SentenceTransformers.load_model(repo) + assert {:ok, tokenizer} = SentenceTransformers.load_tokenizer(repo) + + serving = SentenceTransformers.text_embedding(model_info, tokenizer) + result = Nx.Serving.run(serving, "query: Cats are cute.") + + assert %{embedding: embedding} = result + assert Nx.shape(embedding) == {384} + + # Reference values from PyTorch sentence-transformers 6.1.0 for "query: Cats are cute.": + # [-0.010973, 0.073814, 0.011410] + slice = Nx.to_flat_list(embedding[1..3]) + py_reference = [-0.01097320, 0.07381421, 0.01141022] + + slice + |> Enum.zip(py_reference) + |> Enum.each(fn {val, ref} -> + assert_in_delta val, ref, 1.0e-5 + end) + end end