similarity

class sentence_transformers.util.similarity.SimilarityFunction(value)[source]

Enum class for supported similarity functions. The following functions are supported:

  • SimilarityFunction.COSINE ("cosine"): Cosine similarity

  • SimilarityFunction.DOT_PRODUCT ("dot", dot_product): Dot product similarity

  • SimilarityFunction.EUCLIDEAN ("euclidean"): Euclidean distance

  • SimilarityFunction.MANHATTAN ("manhattan"): Manhattan distance

  • SimilarityFunction.MAXSIM ("maxsim"): Late-interaction MaxSim, used by MultiVectorEncoder (ColBERT-style) models.

  • SimilarityFunction.MEAN_MAXSIM ("meanmaxsim"): MaxSim divided by the query’s token count, so scores are comparable across query lengths. Same rankings within a query.

static possible_values() → list[str][source]

Returns a list of possible values for the SimilarityFunction enum.

Returns:

A list of possible values for the SimilarityFunction enum.

Return type:

list

Example

>>> possible_values = SimilarityFunction.possible_values()
>>> possible_values
['cosine', 'dot', 'euclidean', 'manhattan', 'maxsim', 'meanmaxsim']
static to_similarity_fn(similarity_function: str | SimilarityFunction) → Callable[[list | ndarray | Tensor, list | ndarray | Tensor], Tensor][source]

Converts a similarity function name or enum value to the corresponding similarity function.

Parameters:

similarity_function (Union[str, SimilarityFunction]) – The name or enum value of the similarity function.

Returns:

The corresponding similarity function. The MaxSim family also accepts further keyword arguments, see maxsim().

Return type:

Callable[[Union[list, np.ndarray, Tensor], Union[list, np.ndarray, Tensor]], Tensor]

Raises:

ValueError – If the provided function is not supported.

Example

>>> similarity_fn = SimilarityFunction.to_similarity_fn("cosine")
>>> similarity_scores = similarity_fn(embeddings1, embeddings2)
>>> similarity_scores
tensor([[0.3952, 0.0554],
        [0.0992, 0.1570]])
static to_similarity_pairwise_fn(similarity_function: str | SimilarityFunction) → Callable[[list | ndarray | Tensor, list | ndarray | Tensor], Tensor][source]

Converts a similarity function into a pairwise similarity function.

The pairwise similarity function returns the diagonal vector from the similarity matrix, i.e. it only computes the similarity(a[i], b[i]) for each i in the range of the input tensors, rather than computing the similarity between all pairs of a and b.

Parameters:

similarity_function (Union[str, SimilarityFunction]) – The name or enum value of the similarity function.

Returns:

The pairwise similarity function. The MaxSim family also accepts further keyword arguments, see maxsim_pairwise().

Return type:

Callable[[Union[list, np.ndarray, Tensor], Union[list, np.ndarray, Tensor]], Tensor]

Raises:

ValueError – If the provided similarity function is not supported.

Example

>>> pairwise_fn = SimilarityFunction.to_similarity_pairwise_fn("cosine")
>>> similarity_scores = pairwise_fn(embeddings1, embeddings2)
>>> similarity_scores
tensor([0.3952, 0.1570])
sentence_transformers.util.similarity.cos_sim(a: list | ndarray | Tensor, b: list | ndarray | Tensor) → Tensor[source]

Computes the cosine similarity between two tensors.

Parameters:

  • a (Union[list, np.ndarray, Tensor]) – The first tensor.

  • b (Union[list, np.ndarray, Tensor]) – The second tensor.

Returns:

Matrix with res[i][j] = cos_sim(a[i], b[j])

Return type:

Tensor

sentence_transformers.util.similarity.dot_score(a: list | ndarray | Tensor, b: list | ndarray | Tensor) → Tensor[source]

Computes the dot-product dot_prod(a[i], b[j]) for all i and j.

Parameters:

  • a (Union[list, np.ndarray, Tensor]) – The first tensor.

  • b (Union[list, np.ndarray, Tensor]) – The second tensor.

Returns:

Matrix with res[i][j] = dot_prod(a[i], b[j])

Return type:

Tensor

sentence_transformers.util.similarity.euclidean_sim(a: list | ndarray | Tensor, b: list | ndarray | Tensor) → Tensor[source]

Computes the euclidean similarity (i.e., negative distance) between two tensors. Handles sparse tensors without converting to dense when possible.

Parameters:

  • a (Union[list, np.ndarray, Tensor]) – The first tensor.

  • b (Union[list, np.ndarray, Tensor]) – The second tensor.

Returns:

Matrix with res[i][j] = -euclidean_distance(a[i], b[j])

Return type:

Tensor

sentence_transformers.util.similarity.manhattan_sim(a: list | ndarray | Tensor, b: list | ndarray | Tensor) → Tensor[source]

Computes the manhattan similarity (i.e., negative distance) between two tensors. Handles sparse tensors without converting to dense when possible.

Parameters:

  • a (Union[list, np.ndarray, Tensor]) – The first tensor.

  • b (Union[list, np.ndarray, Tensor]) – The second tensor.

Returns:

Matrix with res[i][j] = -manhattan_distance(a[i], b[j])

Return type:

Tensor

sentence_transformers.util.similarity.maxsim(a: list | ndarray | Tensor, b: list | ndarray | Tensor, a_mask: Tensor | None = None, b_mask: Tensor | None = None, chunk_elements: int | None = None, length_normalize: bool = False, device: str | device | None = None) → Tensor[source]

Computes the MaxSim (late-interaction) score between two collections of multi-vector embeddings.

For each query in a and document in b, the score is the sum over query tokens of the maximum similarity to any document token: sum_i max_j (a_i . b_j). This is the scoring function used by ColBERT-style models.

Parameters:

  • a (Union[list, np.ndarray, Tensor]) – Query embeddings. Either a 3D tensor of shape (batch_a, num_query_tokens, embedding_dim) (pre-padded), a list of 2D tensors of shape (num_query_tokens_i, embedding_dim) (variable-length, padded internally), or a single 2D tensor or array (one query, e.g. from a singular encode call, scored as a batch of one).

  • b (Union[list, np.ndarray, Tensor]) – Document embeddings. Either a 3D tensor of shape (batch_b, num_doc_tokens, embedding_dim), a list of 2D tensors of shape (num_doc_tokens_i, embedding_dim), or a single 2D tensor or array (one document, scored as a batch of one).

  • a_mask (Tensor, optional) – Boolean or float mask for query tokens, shape (batch_a, num_query_tokens), (num_query_tokens,) alongside a single 2D a, or a list of per-item 1D masks (padded to the batch). Tokens with a 0 / False entry are excluded from the sum. Defaults to None (use all tokens).

  • b_mask (Tensor, optional) – Boolean or float mask for document tokens, shape (batch_b, num_doc_tokens), (num_doc_tokens,) alongside a single 2D b, or a list of per-item 1D masks. Tokens with a 0 / False entry are excluded from the max. Defaults to None (use all tokens).

  • chunk_elements (int, optional) – Element budget for the padded (chunk, d_tokens, dim) documents plus the 4D (batch_a, chunk, q_tokens, d_tokens) scoring intermediate. Documents are greedily packed into chunks (padded per chunk, so a long outlier only sizes its own chunk) that stay under the budget, with a floor of one document per chunk. Defaults to None, which applies a 100M-element budget (at most ~400 MB, half that in bf16 / fp16). With very large query batches that floor can still be big: shard queries externally if a single document exceeds memory.

  • length_normalize (bool, optional) – Divide each score by the number of real (unmasked) query tokens, yielding MeanMaxSim: scores land in the per-token similarity range (about [-1, 1] for normalized embeddings) independent of the query length. Defaults to False.

  • device (str, torch.device, optional) – Device to run the scoring on, moving each budget-sized document chunk there rather than the whole corpus. The returned scores stay on the documents’ device either way. Defaults to None (the documents’ device).

Returns:

Matrix with res[i][j] = MaxSim(a[i], b[j]), shape (batch_a, batch_b), always float32 (half precision inputs are accumulated in float32 to keep nearby scores distinct), on the documents’ device: CPU documents score on the CPU even against GPU queries.

Return type:

Tensor

sentence_transformers.util.similarity.maxsim_pairwise(a: list | ndarray | Tensor, b: list | ndarray | Tensor, a_mask: Tensor | None = None, b_mask: Tensor | None = None, chunk_elements: int | None = None, length_normalize: bool = False, device: str | device | None = None) → Tensor[source]

Computes the pairwise MaxSim (late-interaction) score between each query-document pair.

For each i, computes the MaxSim score between a[i] and b[i]: the sum over query tokens of the maximum similarity to any document token. This is the pairwise analogue of maxsim().

Parameters:

  • a (Union[list, np.ndarray, Tensor]) – Query embeddings. Either a 3D tensor of shape (batch, num_query_tokens, embedding_dim) (pre-padded), a list of 2D tensors with shape (num_query_tokens_i, embedding_dim), or a single 2D tensor or array (one query, e.g. from a singular encode call, scored as a batch of one).

  • b (Union[list, np.ndarray, Tensor]) – Document embeddings. Either a 3D tensor of shape (batch, num_doc_tokens, embedding_dim), a list of 2D tensors with shape (num_doc_tokens_i, embedding_dim), or a single 2D tensor or array (one document, scored as a batch of one).

  • a_mask (Tensor, optional) – Boolean or float mask for query tokens, shape (batch, num_query_tokens), (num_query_tokens,) alongside a single 2D a, or a list of per-item 1D masks (padded to the batch). Tokens with a 0 / False entry are excluded from the sum. Defaults to None (use all tokens).

  • b_mask (Tensor, optional) – Boolean or float mask for document tokens, shape (batch, num_doc_tokens), (num_doc_tokens,) alongside a single 2D b, or a list of per-item 1D masks. Tokens with a 0 / False entry are excluded from the max. Defaults to None (use all tokens).

  • chunk_elements (int, optional) – Element budget for the padded (chunk, q_tokens, dim) queries and (chunk, d_tokens, dim) documents plus the (chunk, q_tokens, d_tokens) scoring intermediate. Pairs are greedily packed into chunks (padded per chunk, so a long outlier only sizes its own chunk) that stay under the budget, with a floor of one pair per chunk. Defaults to None, which applies a 100M-element budget (at most ~400 MB, half that in bf16 / fp16).

  • length_normalize (bool, optional) – Divide each score by the number of real (unmasked) query tokens, yielding MeanMaxSim: scores land in the per-token similarity range (about [-1, 1] for normalized embeddings) independent of the query length. Defaults to False.

  • device (str, torch.device, optional) – Device to run the scoring on, moving each budget-sized chunk of pairs there rather than everything. The returned scores stay on the documents’ device either way. Defaults to None (the documents’ device).

Returns:

Vector with res[i] = MaxSim(a[i], b[i]), shape (batch,), always float32 (half precision inputs are accumulated in float32 to keep nearby scores distinct), on the documents’ device: CPU documents score on the CPU even against GPU queries.

Return type:

Tensor

sentence_transformers.util.similarity.mean_maxsim(a: list | ndarray | Tensor, b: list | ndarray | Tensor, a_mask: Tensor | None = None, b_mask: Tensor | None = None, chunk_elements: int | None = None, length_normalize: bool = True, device: str | device | None = None) → Tensor[source]

Computes the MeanMaxSim score between two collections of multi-vector embeddings: maxsim() divided by each query’s real token count, with maxsim()’s arguments and return shape.

MaxSim sums one similarity per query token, so its scale grows with query length and scores are not comparable across queries. MeanMaxSim averages instead, landing in the per-token similarity range (about [-1, 1] for normalized embeddings). Rankings within a query are unchanged, since the divisor is constant per row.

length_normalize defaults to True here (the Mean in the name): False recovers plain maxsim(), so the whole MaxSim family accepts the same keywords.

sentence_transformers.util.similarity.mean_maxsim_pairwise(a: list | ndarray | Tensor, b: list | ndarray | Tensor, a_mask: Tensor | None = None, b_mask: Tensor | None = None, chunk_elements: int | None = None, length_normalize: bool = True, device: str | device | None = None) → Tensor[source]

Computes the pairwise MeanMaxSim score for each query-document pair: maxsim_pairwise() divided by each query’s real token count, with its arguments and return shape. See mean_maxsim() for why. length_normalize defaults to True here, and False recovers plain maxsim_pairwise().

sentence_transformers.util.similarity.pairwise_angle_sim(x: Tensor, y: Tensor) → Tensor[source]

Computes the absolute normalized angle distance. See AnglELoss or https://huggingface.co/papers/2309.12871 for more information.

Parameters:

  • x (Tensor) – The first tensor.

  • y (Tensor) – The second tensor.

Returns:

Vector with res[i] = angle_sim(a[i], b[i])

Return type:

Tensor

sentence_transformers.util.similarity.pairwise_cos_sim(a: list | ndarray | Tensor, b: list | ndarray | Tensor) → Tensor[source]

Computes the pairwise cosine similarity cos_sim(a[i], b[i]).

Parameters:

  • a (Union[list, np.ndarray, Tensor]) – The first tensor.

  • b (Union[list, np.ndarray, Tensor]) – The second tensor.

Returns:

Vector with res[i] = cos_sim(a[i], b[i])

Return type:

Tensor

sentence_transformers.util.similarity.pairwise_dot_score(a: list | ndarray | Tensor, b: list | ndarray | Tensor) → Tensor[source]

Computes the pairwise dot-product dot_prod(a[i], b[i]).

Parameters:

  • a (Union[list, np.ndarray, Tensor]) – The first tensor.

  • b (Union[list, np.ndarray, Tensor]) – The second tensor.

Returns:

Vector with res[i] = dot_prod(a[i], b[i])

Return type:

Tensor

sentence_transformers.util.similarity.pairwise_euclidean_sim(a: list | ndarray | Tensor, b: list | ndarray | Tensor) → Tensor[source]

Computes the euclidean distance (i.e., negative distance) between pairs of tensors.

Parameters:

  • a (Union[list, np.ndarray, Tensor]) – The first tensor.

  • b (Union[list, np.ndarray, Tensor]) – The second tensor.

Returns:

Vector with res[i] = -euclidean_distance(a[i], b[i])

Return type:

Tensor

sentence_transformers.util.similarity.pairwise_manhattan_sim(a: list | ndarray | Tensor, b: list | ndarray | Tensor) → Tensor[source]

Computes the manhattan similarity (i.e., negative distance) between pairs of tensors.

Parameters:

  • a (Union[list, np.ndarray, Tensor]) – The first tensor.

  • b (Union[list, np.ndarray, Tensor]) – The second tensor.

Returns:

Vector with res[i] = -manhattan_distance(a[i], b[i])

Return type:

Tensor

sentence_transformers.util.similarity.pytorch_cos_sim(a: Tensor, b: Tensor) → Tensor[source]

Computes the cosine similarity between two tensors.

Parameters:

  • a (Union[list, np.ndarray, Tensor]) – The first tensor.

  • b (Union[list, np.ndarray, Tensor]) – The second tensor.

Returns:

Matrix with res[i][j] = cos_sim(a[i], b[j])

Return type:

Tensor