tensor

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

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

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

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

Keeps only the top-k values (in absolute terms) for each embedding and creates a sparse tensor.

Parameters:

• embeddings (Union[np.ndarray, torch.Tensor]) – Embeddings to sparsify by keeping only top_k values.

• max_active_dims (int) – Number of values to keep as non-zeros per embedding.

Returns:

A sparse tensor containing only the top-k values per embedding.

Return type:

torch.Tensor

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

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]