unroll colbert implementation

server.py

@@ -1,42 +1,172 @@
-from flask import Flask, render_template, request
+import math, os, ujson, tqdm
+import torch
+import torch.nn.functional as F
+
+from itertools import product
+from flask import Flask, request
 from functools import lru_cache
-import math
-import os
 from dotenv import load_dotenv
 
 from colbert import Searcher
 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()
 
 INDEX_NAME = os.getenv("INDEX_NAME", 'index')
 INDEX_ROOT = os.getenv("INDEX_ROOT", '.')
+INDEX_PATH = os.path.join(INDEX_ROOT, INDEX_NAME)
+
 PORT = int(os.getenv("PORT", 8893))
 app = Flask(__name__)
 
 searcher = Searcher(index_root=INDEX_ROOT, index=INDEX_NAME)
-ranker = IndexScorer(
-    index_path=os.path.join(INDEX_ROOT, INDEX_NAME), 
-    use_gpu=False, 
-    load_index_with_mmap=False
+ranker = IndexScorer(index_path=INDEX_PATH, use_gpu=False, load_index_with_mmap=False)
+
+searcher.configure(
+    ncells=1, 
+    centroid_score_threshold=0.5, 
+    ndocs=256
 )
 
 counter = {"api" : 0}
 
+def init_index(index_path=INDEX_PATH):
+    """
+    Load all tensors necessary for running ColBERT
+    """
+    global centroids, embeddings, ivf, doclens, metadata, bucket_weights, codec, offsets
+    with open(os.path.join(index_path, 'metadata.json')) as f:
+        metadata = ujson.load(f)
+
+    centroids = torch.load(os.path.join(index_path, 'centroids.pt'), map_location='cpu')
+    centroids = centroids.float()
+
+    ivf, ivf_lengths = torch.load(os.path.join(index_path, "ivf.pid.pt"), map_location='cpu')
+    ivf = StridedTensor(ivf, ivf_lengths, use_gpu=False)
+
+    embeddings = ResidualCodec.Embeddings.load_chunks(
+        index_path,
+        range(metadata['num_chunks']),
+        metadata['num_embeddings'],
+        load_index_with_mmap=False,
+    )
+
+    doclens = []
+    for chunk_idx in tqdm.tqdm(range(metadata['num_chunks'])):
+        with open(os.path.join(index_path, f'doclens.{chunk_idx}.json')) as f:
+            chunk_doclens = ujson.load(f)
+            doclens.extend(chunk_doclens)
+    doclens = torch.tensor(doclens)
+
+    buckets_path = os.path.join(index_path, 'buckets.pt')
+    bucket_cutoffs, bucket_weights = torch.load(buckets_path, map_location='cpu')
+    bucket_weights = bucket_weights.float()
+
+    codec = ResidualCodec.load(index_path)
+
+    load_index_with_mmap = False
+    if load_index_with_mmap:
+        assert metadata['num_chunks'] == 1
+        offsets = torch.cumsum(doclens, dim=0)
+        offsets = torch.cat( (torch.zeros(1, dtype=torch.int64), offsets) )
+    else:
+        embeddings_strided = ResidualEmbeddingsStrided(codec, embeddings, doclens)
+        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(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_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
+
+    return ColBERT.segmented_maxsim(scores, D_lengths)
+
+
+def score_pids(config, Q, pids, centroid_scores):
+    idx = centroid_scores.max(-1).values >= config.centroid_score_threshold
+
+    pids = IndexScorer.filter_pids(
+        pids, centroid_scores, embeddings.codes, doclens,
+        offsets, idx, config.ndocs
+    )
+
+    # Rank final list of docs using full approximate embeddings (including residuals)
+    D_packed = IndexScorer.decompress_residuals(
+        pids, doclens, offsets, bucket_weights, codec.reversed_bit_map,
+        codec.decompression_lookup_table, embeddings.residuals, embeddings.codes, 
+        centroids, codec.dim, metadata['config']['nbits']
+    )
+    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, config)
+
+    return scores, pids
+
+    
+def generate_candidates(Q):
+    ncells = searcher.config.ncells
+    Q = Q.squeeze(0)
+
+    # Get the closest centroids via a matrix multiplication + argmax
+    centroid_scores = (centroids @ Q.T)
+    if ncells == 1:
+        cells = centroid_scores.argmax(dim=0, keepdim=True).permute(1, 0)
+    else:
+        cells = centroid_scores.topk(ncells, dim=0, sorted=False).indices.permute(1, 0)  # (32, ncells)
+    cells = cells.flatten().contiguous()  # (32 * ncells,)
+    cells = cells.unique(sorted=False)
+
+    # (?) Find the relevant passages related to each cluster
+    pids, _ = ivf.lookup(cells)
+
+    # Sort and retun values
+    sorter = pids.sort()
+    pids = sorter.values
+    pids, _ = torch.unique_consecutive(pids, return_counts=True)
+    return pids, centroid_scores
+
 
 def search_colbert(query, k):
     # Add the appropriate [Q], [D] tokens and encode with ColBERT
-    searcher.configure(ncells=1, centroid_score_threshold=0.5, ndocs=256)
     Q = searcher.encode(query)
 
     # Cut off query to maxlen tokens
     Q = Q[:, :searcher.config.query_maxlen]
 
     # Find the passage candidates (i.e., closest candidates to the Q centroid)
-    pids, centroid_scores = ranker.generate_candidates(searcher.config, Q)
+    pids, centroid_scores = generate_candidates(Q)
 
     # Use our index to calculate the max similarity scores
-    scores, pids = ranker.score_pids(searcher.config, Q, pids, centroid_scores)
+    scores, pids = score_pids(searcher.config, Q, pids, centroid_scores)
 
     # Sort and return values
     scores_sorter = scores.sort(descending=True)
@@ -51,13 +181,17 @@ def search_colbert(query, k):
 @lru_cache(maxsize=1000000)
 def api_search_query(query, k):
     print(f"Query={query}")
+
     k = 10 if k == None else min(int(k), 100)
 
+    # Use ColBERT to find passages related to the query
     pids, ranks, scores = search_colbert(query, k)
 
+    # Softmax output probs
     probs = [math.exp(s) for s in scores]
     probs = [p / sum(probs) for p in probs]
 
+    # Compile and return using the API
     topk = []
     for pid, rank, score, prob in zip(pids, ranks, scores, probs):
         text = searcher.collection[pid]
@@ -86,9 +220,10 @@ def api_search():
 if __name__ == "__main__":
     """
     Example usage:
-    INDEX_ROOT=/Users/dhei/personal/4440/project/colbert-acl INDEX_NAME=index python server.py
+    python server.py
     http://localhost:8893/api/search?k=25&query=How to extend context windows?
     """
-    print(api_search_query("This is a test", 1))
+    init_index(index_path=INDEX_PATH)
+    print(api_search_query("This is a test", 2))
     # app.run("0.0.0.0", PORT)