TorchSparseIndex

class lightning_ir.retrieve.pytorch.sparse_searcher.TorchSparseIndex(index_dir: Path, similarity_function: 'dot' | 'cosine', use_gpu: bool = False)[source]

Bases: object

Torch-based sparse index for efficient retrieval.

__init__(index_dir: Path, similarity_function: 'dot' | 'cosine', use_gpu: bool = False) → None[source]

Initialize the TorchSparseIndex.

Parameters:

  • index_dir (Path) – Directory containing the index files.

  • similarity_function (Literal["dot", "cosine"]) – The similarity function to use.

  • use_gpu (bool) – Whether to use GPU for computations. Defaults to False.

Raises:

ValueError – If the similarity function is not recognized.

Methods

__init__(index_dir, similarity_function[, ...])

Initialize the TorchSparseIndex.

cosine_similarity(x, y)

Compute cosine similarity between two tensors.

dot_similarity(x, y)

Compute dot product similarity between two tensors.

score(embeddings)

Compute scores for the given embeddings.

to_gpu()

Move the index to GPU if available.

Attributes

num_embeddings

Get the number of embeddings in the index.
static cosine_similarity(x: Tensor, y: Tensor) → Tensor[source]

Compute cosine similarity between two tensors.

Parameters:

  • x (torch.Tensor) – The first tensor.

  • y (torch.Tensor) – The second tensor.

Returns:

The cosine similarity scores.

Return type:

torch.Tensor

static dot_similarity(x: Tensor, y: Tensor) → Tensor[source]

Compute dot product similarity between two tensors.

Parameters:

  • x (torch.Tensor) – The first tensor.

  • y (torch.Tensor) – The second tensor.

Returns:

The dot product similarity scores.

Return type:

torch.Tensor

property num_embeddings: int

Get the number of embeddings in the index.

Returns:

The number of embeddings.

Return type:

int

score(embeddings: Tensor) → Tensor[source]

Compute scores for the given embeddings.

Parameters:

embeddings (torch.Tensor) – The embeddings to score.

Returns:

The computed scores.

Return type:

torch.Tensor

to_gpu() → None[source]

Move the index to GPU if available.