remove colbert dependency

search.py

@@ -9,7 +9,6 @@ from colbert.search.index_storage import IndexScorer
 from colbert.search.strided_tensor import StridedTensor
 from colbert.indexing.codecs.residual_embeddings_strided import ResidualEmbeddingsStrided
 from colbert.indexing.codecs.residual import ResidualCodec
-from colbert.modeling.colbert import ColBERT
 
 load_dotenv()
 
@@ -67,19 +66,23 @@ def init_colbert(index_path=INDEX_PATH, load_index_with_mmap=False):
         offsets = embeddings_strided.codes_strided.offsets
 
     
-# def colbert_score_reduce(scores_padded, D_mask):
-#     D_padding = ~D_mask.view(scores_padded.size(0), scores_padded.size(1)).bool()
-#     scores_padded[D_padding] = -9999
-#     scores = scores_padded.max(1).values
-#     return scores.sum(-1)
+def colbert_score_reduce(scores_padded, D_mask):
+    D_padding = ~D_mask.view(scores_padded.size(0), scores_padded.size(1)).bool()
+    scores_padded[D_padding] = -9999
+    scores = scores_padded.max(1).values
 
+    return scores.sum(-1)
 
-# def colbert_score(Q, D_padded, D_mask):
-#     assert Q.dim() == 3, Q.size()
-#     assert D_padded.dim() == 3, D_padded.size()
-#     assert Q.size(0) in [1, D_padded.size(0)]
-#     scores = D_padded @ Q.to(dtype=D_padded.dtype).permute(0, 2, 1)
-#     return colbert_score_reduce(scores, D_mask)
+
+def colbert_score(Q, D_padded, D_mask):
+    assert Q.dim() == 3, Q.size()
+    assert D_padded.dim() == 3, D_padded.size()
+    assert Q.size(0) in [1, D_padded.size(0)]
+
+    scores = D_padded @ Q.to(dtype=D_padded.dtype).permute(0, 2, 1)
+    scores = colbert_score_reduce(scores, D_mask)
+
+    return scores
 
 
 def colbert_score_packed(Q, D_packed, D_lengths):
@@ -91,8 +94,8 @@ def colbert_score_packed(Q, D_packed, D_lengths):
 
     scores = D_packed @ Q.T
 
-    # C++ : Calculate maxsim operation
-    scores = ColBERT.segmented_maxsim(scores, D_lengths)
+    scores_padded, scores_mask = StridedTensor(scores, D_lengths, use_gpu=False).as_padded_tensor()
+    scores = colbert_score_reduce(scores_padded, scores_mask)
 
     return scores
 
@@ -116,10 +119,10 @@ def score_pids(config, Q, pids, centroid_scores):
 
     if Q.size(0) == 1:
         scores = colbert_score_packed(Q, D_packed, D_mask)
-    # else:
-    #     D_strided = StridedTensor(D_packed, D_mask, use_gpu=False)
-    #     D_padded, D_lengths = D_strided.as_padded_tensor()
-    #     scores = colbert_score(Q, D_padded, D_lengths, config)
+    else:
+        D_strided = StridedTensor(D_packed, D_mask, use_gpu=False)
+        D_padded, D_lengths = D_strided.as_padded_tensor()
+        scores = colbert_score(Q, D_padded, D_lengths)
 
     return scores, pids