similarity

class sentence_transformers.util.similarity.SimilarityFunction(value)

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.

static possible_values() → list[str]

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']
static to_similarity_fn(similarity_function: str | SimilarityFunction) → Callable[[Tensor | ndarray, Tensor | ndarray], Tensor]

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.

Return type:

Callable[[Union[Tensor, ndarray], Union[Tensor, ndarray]], 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[[Tensor | ndarray, Tensor | ndarray], Tensor]

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.

Return type:

Callable[[Union[Tensor, ndarray], Union[Tensor, ndarray]], 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

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

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

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

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, document_chunk_size: int | None = None) → Tensor

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) or a list of 2D tensors of shape (num_query_tokens_i, embedding_dim) (variable-length, padded internally).

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

  • a_mask (Tensor, optional) – Boolean or float mask for query tokens, shape (batch_a, num_query_tokens). 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). Tokens with a 0 / False entry are excluded from the max. Defaults to None (use all tokens).

  • document_chunk_size (int, optional) – If set, iterate the einsum + max-reduction over document chunks of this size along the b axis. Keeps the full (a, b, s, t) 4D intermediate from being materialized at once. Useful for evaluation against large corpora. Defaults to None (single einsum over the full b axis).

Returns:

Matrix with res[i][j] = MaxSim(a[i], b[j]), shape (batch_a, batch_b).

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) → Tensor

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) or a list of 2D tensors with shape (num_query_tokens_i, embedding_dim).

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

  • a_mask (Tensor, optional) – Boolean or float mask for query tokens, shape (batch, num_query_tokens). 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). Tokens with a 0 / False entry are excluded from the max. Defaults to None (use all tokens).

Returns:

Vector with res[i] = MaxSim(a[i], b[i]), shape (batch,).

Return type:

Tensor

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

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: Tensor, b: Tensor) → Tensor

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: Tensor, b: Tensor) → Tensor

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

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

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

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