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.