fix search path bug

.gitignore

@@ -1 +1,2 @@
-__pycache__
+__pycache__
+experiments

search.py

@@ -1,10 +1,11 @@
-import os, ujson, tqdm
+import os, shutil, ujson, tqdm
 import torch
 import torch.nn.functional as F
 
 from dotenv import load_dotenv
 
 from colbert import Searcher
+from colbert.infra.config.settings import RunSettings
 from colbert.search.index_storage import IndexScorer
 from colbert.search.strided_tensor import StridedTensor
 from colbert.indexing.codecs.residual_embeddings_strided import ResidualEmbeddingsStrided
@@ -15,8 +16,16 @@ load_dotenv()
 INDEX_NAME = os.getenv("INDEX_NAME", 'index')
 INDEX_ROOT = os.getenv("INDEX_ROOT", '.')
 INDEX_PATH = os.path.join(INDEX_ROOT, INDEX_NAME)
-searcher = Searcher(index_root=INDEX_ROOT, index=INDEX_NAME)
 
+# Move index to ColBERT experiment path
+src_path = os.path.join(INDEX_ROOT, INDEX_NAME)
+dest_path = os.path.join(INDEX_ROOT, 'experiments', 'default', 'indexes', INDEX_NAME)
+if not os.path.exists(dest_path):
+    print(f'Copying {src_path} -> {dest_path}')
+    os.makedirs(dest_path)
+    shutil.copytree(src_path, dest_path, dirs_exist_ok=True)
+
+searcher = Searcher(index=INDEX_NAME)
 searcher.configure(
     ncells=1, 
     centroid_score_threshold=0.5, 
@@ -65,13 +74,25 @@ def init_colbert(index_path=INDEX_PATH, load_index_with_mmap=False):
         embeddings_strided = ResidualEmbeddingsStrided(codec, embeddings, doclens)
         offsets = embeddings_strided.codes_strided.offsets
 
+
+# def colbert_score_packed(Q, D_packed, D_lengths):
+#     Q = Q.squeeze(0)
+#     Q = Q.to(dtype=D_packed.dtype)
+#     assert Q.dim() == 2, Q.size()
+#     assert D_packed.dim() == 2, D_packed.size()
+#     scores = D_packed @ Q.T
+#     scores_padded, scores_mask = StridedTensor(scores, D_lengths, use_gpu=False).as_padded_tensor()
+#     scores = colbert_score_reduce(scores_padded, scores_mask)
+#     return scores
+
     
 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
+    scores = scores.sum(-1)
 
-    return scores.sum(-1)
+    return scores
 
 
 def colbert_score(Q, D_padded, D_mask):
@@ -85,21 +106,6 @@ def colbert_score(Q, D_padded, D_mask):
     return scores
 
 
-def colbert_score_packed(Q, D_packed, D_lengths):
-    Q = Q.squeeze(0)
-    Q = Q.to(dtype=D_packed.dtype)
-
-    assert Q.dim() == 2, Q.size()
-    assert D_packed.dim() == 2, D_packed.size()
-
-    scores = D_packed @ Q.T
-
-    scores_padded, scores_mask = StridedTensor(scores, D_lengths, use_gpu=False).as_padded_tensor()
-    scores = colbert_score_reduce(scores_padded, scores_mask)
-
-    return scores
-
-
 def score_pids(config, Q, pids, centroid_scores):
     # C++ : Filter pids under the centroid score threshold
     idx = centroid_scores.max(-1).values >= config.centroid_score_threshold
@@ -117,12 +123,12 @@ def score_pids(config, Q, pids, centroid_scores):
     D_packed = F.normalize(D_packed.to(torch.float32), p=2, dim=-1)
     D_mask = doclens[pids.long()]
 
-    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)
+    # 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)
 
     return scores, pids
 

server.py

@@ -68,5 +68,5 @@ if __name__ == "__main__":
     """
     init_colbert()
     # print(api_search_query("This is a test", 2))
+    print(f'Test it at: http://localhost:8893/api/search?k=25&query=How to extend context windows?')
     app.run("0.0.0.0", PORT)
-