Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added a note to the documentation stating that the similarity method does not support embeddings other than non-quantized ones. #3131

Merged
merged 1 commit into from
Dec 12, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 11 additions & 5 deletions sentence_transformers/SentenceTransformer.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,11 @@
from typing import Any, Callable, Literal, overload

import numpy as np
import numpy.typing as npt
import torch
import torch.multiprocessing as mp
import transformers
from huggingface_hub import HfApi
from numpy import ndarray
from torch import Tensor, device, nn
from tqdm.autonotebook import trange
from transformers import is_torch_npu_available
Expand Down Expand Up @@ -723,14 +723,15 @@ def similarity_fn_name(
def similarity(self, embeddings1: Tensor, embeddings2: Tensor) -> Tensor: ...

@overload
def similarity(self, embeddings1: ndarray, embeddings2: ndarray) -> Tensor: ...
def similarity(self, embeddings1: npt.NDArray[np.float32], embeddings2: npt.NDArray[np.float32]) -> Tensor: ...
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice, looks to be the correct usage. I wasn't aware of these.
https://numpy.org/devdocs/reference/typing.html#numpy.typing.NDArray


@property
def similarity(self) -> Callable[[Tensor | ndarray, Tensor | ndarray], Tensor]:
def similarity(self) -> Callable[[Tensor | npt.NDArray[np.float32], Tensor | npt.NDArray[np.float32]], Tensor]:
"""
Compute the similarity between two collections of embeddings. The output will be a matrix with the similarity
scores between all embeddings from the first parameter and all embeddings from the second parameter. This
differs from `similarity_pairwise` which computes the similarity between each pair of embeddings.
This method supports only embeddings with fp32 precision and does not accommodate quantized embeddings.

Args:
embeddings1 (Union[Tensor, ndarray]): [num_embeddings_1, embedding_dim] or [embedding_dim]-shaped numpy array or torch tensor.
Expand Down Expand Up @@ -772,13 +773,18 @@ def similarity(self) -> Callable[[Tensor | ndarray, Tensor | ndarray], Tensor]:
def similarity_pairwise(self, embeddings1: Tensor, embeddings2: Tensor) -> Tensor: ...

@overload
def similarity_pairwise(self, embeddings1: ndarray, embeddings2: ndarray) -> Tensor: ...
def similarity_pairwise(
self, embeddings1: npt.NDArray[np.float32], embeddings2: npt.NDArray[np.float32]
) -> Tensor: ...

@property
def similarity_pairwise(self) -> Callable[[Tensor | ndarray, Tensor | ndarray], Tensor]:
def similarity_pairwise(
self,
) -> Callable[[Tensor | npt.NDArray[np.float32], Tensor | npt.NDArray[np.float32]], Tensor]:
"""
Compute the similarity between two collections of embeddings. The output will be a vector with the similarity
scores between each pair of embeddings.
This method supports only embeddings with fp32 precision and does not accommodate quantized embeddings.

Args:
embeddings1 (Union[Tensor, ndarray]): [num_embeddings, embedding_dim] or [embedding_dim]-shaped numpy array or torch tensor.
Expand Down
Loading