Skip to content

colbert_kd_scores_top_p

Computes the ColBERT scores between queries and documents embeddings using top-p sampling. This scoring function is dedicated to the knowledge distillation pipeline.

Parameters

  • queries_embeddings (torch.Tensor)

    The queries embeddings. Shape: (batch_size, num_tokens_queries, embedding_size)

  • documents_embeddings (torch.Tensor)

    The documents embeddings. Shape: (batch_size, num_documents, num_tokens_documents, embedding_size)

  • mask (torch.Tensor) – defaults to None

    Optional mask tensor. Shape: (batch_size, num_documents, num_tokens_documents). Default is None.

  • p (float) – defaults to 0.002

    The cumulative probability threshold for top-p sampling. Default is 0.9.