Skip to content

colbert_kd_scores

Computes the ColBERT scores between queries and documents embeddings. This scoring function is dedicated to the knowledge distillation pipeline.

Parameters

  • queries_embeddings ('list | np.ndarray | torch.Tensor')

    The first tensor. The queries embeddings. Shape: (batch_size, num tokens queries, embedding_size)

  • documents_embeddings ('list | np.ndarray | torch.Tensor')

    The second tensor. The documents embeddings. Shape: (batch_size, num candidate documents, num tokens documents, embedding_size)

  • queries_mask ('torch.Tensor') – defaults to None

    The mask for the queries embeddings. Shape: (batch_size, num tokens queries)

  • documents_mask ('torch.Tensor') – defaults to None

    The mask for the documents embeddings. Shape: (batch_size, num candidate documents, num tokens documents)

  • backend ('str | None') – defaults to None

    Scoring backend. See :func:colbert_scores.

Examples

>>> import torch

>>> queries_embeddings = torch.tensor([
...     [[1.], [0.], [0.], [0.]],
...     [[0.], [2.], [0.], [0.]],
...     [[0.], [0.], [3.], [0.]],
... ])

>>> documents_embeddings = torch.tensor([
...     [[[10.], [0.], [1.]], [[20.], [0.], [1.]], [[30.], [0.], [1.]]],
...     [[[0.], [100.], [1.]], [[0.], [200.], [1.]], [[0.], [300.], [1.]]],
...     [[[1.], [0.], [1000.]], [[1.], [0.], [2000.]], [[10.], [0.], [3000.]]],
... ])
>>> documents_mask = torch.tensor([
...     [[0., 1., 1.], [1., 1., 1.], [1., 1., 1.]],
...     [[1., 1., 1.], [1., 1., 1.], [1., 1., 1.]],
...     [[1., 1., 1.], [1., 1., 1.], [1., 1., 1.]],
... ])
>>> query_mask = torch.tensor([
...     [1., 1., 1., 1.], [1., 1., 1., 1.], [1., 1., 0., 1.]
... ])
>>> colbert_kd_scores(
...     queries_embeddings=queries_embeddings,
...     documents_embeddings=documents_embeddings,
...     queries_mask=query_mask,
...     documents_mask=documents_mask,
... )
tensor([[ 1.,  20.,  30.],
        [200., 400., 600.],
        [  0.,   0.,   0.]])