primeqa.ir.dense.colbert_top.colbert.modeling.colbert.colbert_score

primeqa.ir.dense.colbert_top.colbert.modeling.colbert.colbert_score(Q, D_padded, D_mask, config=ColBERTConfig(ncells=None, centroid_score_threshold=None, ndocs=None, index_path=None, index_location=None, nbits=1, kmeans_niters=20, num_partitions_max=10000000, similarity='cosine', bsize=32, accumsteps=1, lr=3e-06, maxsteps=500000, save_every=None, resume=False, resume_optimizer=False, warmup=None, warmup_bert=None, relu=False, nway=2, use_ib_negatives=False, reranker=False, distillation_alpha=1.0, ignore_scores=False, shuffle_every_epoch=False, save_steps=2000, save_epochs=- 1, epochs=10, input_arguments={}, local_models_repository=None, ranks_fn=None, output_dir=None, topK=100, student_teacher_temperature=1.0, student_teacher_top_loss_weight=0.5, teacher_doc_maxlen=180, distill_query_passage_separately=False, query_only=False, loss_function=None, query_weight=0.5, rng_seed=12345, query_maxlen=32, attend_to_mask_tokens=False, interaction='colbert', dim=128, doc_maxlen=180, mask_punctuation=True, checkpoint=None, teacher_checkpoint=None, triples=None, teacher_triples=None, collection=None, queries=None, index_name=None, overwrite=False, root='/home/runner/work/primeqa/primeqa/docs/experiments', experiment='default', index_root=None, name='2024-08/29/18.14.15', rank=0, nranks=1, amp=True, gpus=0))

Supply sizes Q = (1 | num_docs, , dim) and D = (num_docs, *, dim). If Q.size(0) is 1, the matrix will be compared with all passages. Otherwise, each query matrix will be compared against the *aligned passage.

EVENTUALLY: Consider masking with -inf for the maxsim (or enforcing a ReLU).