diff --git a/lib/bumblebee/sentence_transformers.ex b/lib/bumblebee/sentence_transformers.ex new file mode 100644 index 00000000..ca9ec1b1 --- /dev/null +++ b/lib/bumblebee/sentence_transformers.ex @@ -0,0 +1,829 @@ +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 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} -> + {: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), + 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 + + @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 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, + "expected the first module in modules.json to be Transformer or StaticEmbedding"} + end + + {:missing_modules, _reason} -> + {:error, "could not find modules.json in the repository"} + + {:error, reason} -> + {:error, reason} + 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 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} -> + {:error, reason} + 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 -> + 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) + + 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 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 {Keyword.fetch(opts, :prompt), Keyword.fetch(opts, :prompt_name)} do + {{:ok, prompt}, :error} when is_binary(prompt) -> + prompt + + {{:ok, prompt}, :error} when prompt in [false, nil] -> + "" + + {{:ok, prompt}, :error} -> + raise ArgumentError, "expected :prompt to be a string or false, got: #{inspect(prompt)}" + + {:error, {:ok, prompt_name}} when is_binary(prompt_name) -> + find_prompt_by_name!(model_info, prompt_name) + + {:error, {:ok, prompt_name}} when prompt_name in [false, nil] -> + "" + + {:error, {:ok, prompt_name}} -> + raise ArgumentError, + "expected :prompt_name to be a string or false, got: #{inspect(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 + + 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)) + 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 resolve_max_seq_length(transformer, config, tokenizer, root_tokenizer) do + tokenizer_limit = (tokenizer || root_tokenizer || %{})["model_max_length"] + + 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 + + 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 + + @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 + if causal_from_config?(repository) or causal_from_spec?(model_info) do + "last" + else + "mean" + end + 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"] || [] + + 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) + else + _ -> false + end + end + + defp causal_from_spec?(model_info) do + case model_info[:spec] || model_info[:config] do + nil -> + false + + 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 new file mode 100644 index 00000000..a65b25b6 --- /dev/null +++ b/lib/bumblebee/sentence_transformers/downloader.ex @@ -0,0 +1,100 @@ +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 + check_local_file(Path.join(dir, filename)) + end + + def download_file({:local, dir, local_opts}, filename, _opts) when is_binary(dir) do + subdir = local_opts[:subdir] || "" + check_local_file(Path.join([dir, subdir, filename])) + 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 + + 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 + {: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/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 new file mode 100644 index 00000000..b8ab1342 --- /dev/null +++ b/lib/bumblebee/sentence_transformers/pipeline.ex @@ -0,0 +1,280 @@ +defmodule Bumblebee.SentenceTransformers.Pipeline do + @moduledoc false + + alias Axon.MixedPrecision.Policy + alias Bumblebee.Layers + 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. + """ + 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) + + hidden_state = extract_base_hidden_state(base_model, has_weighted_pooling?) + + 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} + + 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(%{ + embedding: embedding, + pooled_state: embedding, + hidden_state: hidden_state + }) + |> apply_type(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 extract_base_hidden_state(base_model, true) do + extract_all_hidden_states(base_model) || extract_single_hidden_state(base_model) + end + + defp extract_base_hidden_state(base_model, false) do + extract_single_hidden_state(base_model) + end + + defp extract_single_hidden_state(base_model) do + 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 + + 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 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 + + _ -> + embedding + end + end + + 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 + Enum.any?(modules, fn m -> + normalize_module_type(m["type"]) == "Normalize" + end) + end + + 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, [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], acc) do + {:ok, Enum.reverse([other | acc])} + end + + defp fuse_dense_modules(_repository, [], acc) do + {:ok, Enum.reverse(acc)} + end + + 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"] + no_bias? = !Map.get(c1, "bias", true) and !Map.get(c2, "bias", true) + dims_match? = c1["out_features"] == c2["in_features"] + + c1_linear? and c2_linear? and no_bias? and dims_match? + end + + defp 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 + + 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_test.exs b/test/bumblebee/sentence_transformers_test.exs new file mode 100644 index 00000000..de8f7714 --- /dev/null +++ b/test/bumblebee/sentence_transformers_test.exs @@ -0,0 +1,1315 @@ +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 + + @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