@spec encode_document(
  map(),
  Bumblebee.Tokenizer.t(),
  String.t() | [String.t()],
  keyword()
) ::
  Nx.Tensor.t() | [Nx.Tensor.t()]

Encodes document text into a rank-2 token matrix, or a list of matrices for a list of inputs. The saved document, passage, or corpus prompt is used unless :prompt or :prompt_name overrides it.

@spec encode_query(
  map(),
  Bumblebee.Tokenizer.t(),
  String.t() | [String.t()],
  keyword()
) ::
  Nx.Tensor.t() | [Nx.Tensor.t()]

Encodes query text into a rank-2 token matrix, or a list of matrices for a list of inputs. The saved query prompt is used unless :prompt or :prompt_name overrides it.

@spec hierarchical_token_pool(Nx.Tensor.t() | [Nx.Tensor.t()], number(), keyword()) ::
  Nx.Tensor.t() | [Nx.Tensor.t()]

Applies HierarchicalTokenPooling to a token matrix or list of token matrices.

This follows SentenceTransformers' saved-module behavior: leading :num_protected_tokens stay untouched and the remaining L2-normalized rows are clustered with Ward linkage on cosine distance before each cluster is replaced by its mean.

@spec load_model(repository(), keyword()) :: {:ok, map()} | {:error, String.t()}

Loads a text multi-vector checkpoint.

The base transformer and its token-level Dense/Normalize modules are built from the saved module manifest. MultiVectorMask configuration is retained on the returned model info and applied by encode_query/4 and encode_document/4. Checkpoints may contain at most one mask module and optional final HierarchicalTokenPooling; multiple masks or other pooling orders are rejected.

@spec maxsim(
  embedding() | [embedding()] | Nx.Tensor.t(),
  embedding() | [embedding()] | Nx.Tensor.t()
) ::
  Nx.Tensor.t()

Computes all-pairs MaxSim scores.

Each input may be a rank-2 token matrix, a rank-3 padded batch, or a list of rank-2 matrices. Padded batches may be supplied as %{embedding: tensor, mask: mask}; masked query tokens contribute zero and masked document tokens are excluded from each maximum.

@spec mean_maxsim(
  embedding() | [embedding()] | Nx.Tensor.t(),
  embedding() | [embedding()] | Nx.Tensor.t()
) :: Nx.Tensor.t()

Computes all-pairs MaxSim scores divided by active query-token count.

@spec similarity(
  embedding() | [embedding()] | Nx.Tensor.t(),
  embedding() | [embedding()] | Nx.Tensor.t(),
  keyword()
) :: Nx.Tensor.t()

Computes MaxSim or MeanMaxSim selected by :similarity_fn_name.

@spec token_embedding(map(), Bumblebee.Tokenizer.t(), keyword()) :: Nx.Serving.t()

Builds a serving that returns rectangular token embeddings before the scoring mask is applied. Each result map uses %{embedding: rank_2_tensor} and includes :tokenized when requested. Prefer encode_query/4 or encode_document/4 for variable-length public token matrices.