tensor

sentence_transformers.util.tensor.batch_to_device(batch: dict[str, Any], target_device: device) → dict[str, Any][source]

Send a PyTorch batch (i.e., a dictionary of string keys to Tensors) to a device (e.g. “cpu”, “cuda”, “mps”).

Parameters:

batch (Dict[str, Tensor]) – The batch to send to the device.
target_device (torch.device) – The target device (e.g. “cpu”, “cuda”, “mps”).

Returns:

The batch with tensors sent to the target device.

Return type:

Dict[str, Tensor]

sentence_transformers.util.tensor.compute_count_vector(embeddings: Tensor) → Tensor[source]

Compute count vector from sparse embeddings indicating how many samples have non-zero values in each dimension.

Parameters:: embeddings – Sparse tensor of shape (batch_size, vocab_size) or (vocab_size,)
Returns:: Count vector of shape (vocab_size,)

sentence_transformers.util.tensor.normalize_embeddings(embeddings: Tensor) → Tensor[source]

Normalizes the embeddings matrix, so that each sentence embedding has unit length.

Parameters:: embeddings (Tensor) – The input embeddings matrix.
Returns:: The normalized embeddings matrix.
Return type:: Tensor

sentence_transformers.util.tensor.select_max_active_dims(embeddings: ndarray | Tensor, max_active_dims: int) → Tensor[source]

sentence_transformers.util.tensor.select_max_active_dims(embeddings: ndarray, max_active_dims: None) → ndarray

sentence_transformers.util.tensor.select_max_active_dims(embeddings: Tensor, max_active_dims: None) → Tensor

Returns a new tensor with only the top-k values (in absolute terms) of each embedding, all others set to zero.

The input embeddings are never modified in place.

Parameters:

embeddings (Union[np.ndarray, torch.Tensor]) – Embeddings to sparsify by keeping only the largest values.
max_active_dims (Optional[int]) – Number of values to keep as non-zeros per embedding. None keeps all values, returning the embeddings as-is.

Raises:

ValueError – If max_active_dims is 0 or negative.

Returns:

A new dense tensor of the same shape, with all but the top-k values per: embedding set to zero. If max_active_dims is None, the embeddings are returned unchanged.

Return type:

Union[np.ndarray, torch.Tensor]

sentence_transformers.util.tensor.to_scipy_coo(x: Tensor) → coo_matrix[source]

Converts a sparse PyTorch tensor to a SciPy COO sparse matrix.

Parameters:: x (torch.Tensor) – A 2-dimensional sparse PyTorch tensor in COO layout.

Example

>>> import torch
>>> from sentence_transformers.util import to_scipy_coo
>>> x = torch.sparse_coo_tensor([[0, 1], [1, 0]], [1.0, 2.0], (3, 3))
>>> to_scipy_coo(x).toarray()
array([[0., 1., 0.],
       [2., 0., 0.],
       [0., 0., 0.]], dtype=float32)

Returns:: A SciPy COO matrix with the same shape, indices and values as the input tensor.
Return type:: scipy.sparse.coo_matrix

sentence_transformers.util.tensor.truncate_embeddings(embeddings: ndarray, truncate_dim: int | None) → ndarray[source]

sentence_transformers.util.tensor.truncate_embeddings(embeddings: Tensor, truncate_dim: int | None) → Tensor

Truncates the embeddings matrix.

Parameters:

embeddings (Union[np.ndarray, torch.Tensor]) – Embeddings to truncate.
truncate_dim (Optional[int]) – The dimension to truncate sentence embeddings to. None does no truncation.

Example

>>> from sentence_transformers import SentenceTransformer
>>> from sentence_transformers.util import truncate_embeddings
>>> model = SentenceTransformer("tomaarsen/mpnet-base-nli-matryoshka")
>>> embeddings = model.encode(["It's so nice outside!", "Today is a beautiful day.", "He drove to work earlier"])
>>> embeddings.shape
(3, 768)
>>> model.similarity(embeddings, embeddings)
tensor([[1.0000, 0.8100, 0.1426],
        [0.8100, 1.0000, 0.2121],
        [0.1426, 0.2121, 1.0000]])
>>> truncated_embeddings = truncate_embeddings(embeddings, 128)
>>> truncated_embeddings.shape
>>> model.similarity(truncated_embeddings, truncated_embeddings)
tensor([[1.0000, 0.8092, 0.1987],
        [0.8092, 1.0000, 0.2716],
        [0.1987, 0.2716, 1.0000]])

Returns:: Truncated embeddings.
Return type:: Union[np.ndarray, torch.Tensor]