colbert local modifications

dsp/modules/colbertv2.py

@@ -2,10 +2,6 @@
 from typing import Any, Optional, Union, List
 
 import requests
-import colbert
-from colbert import Indexer, Searcher
-from colbert.infra import Run, RunConfig, ColBERTConfig
-from colbert.data import Queries, Collection
 from dsp.modules.cache_utils import CacheMemory, NotebookCacheMemory
 from dsp.utils import dotdict
 import os
@@ -81,33 +77,72 @@ def colbertv2_post_request_v2_wrapped(*args, **kwargs):
 os.environ['COLBERT_LOAD_TORCH_EXTENSION_VERBOSE'] = "True"
 
 class ColBERTv2Local:
- def __init__(self,checkpoint:str='colbert-ir/colbertv2.0'):
+ def __init__(self,checkpoint:str='colbert-ir/colbertv2.0',passages:List[str]=[],index_name_or_path:str = "Colbert-RM",experiment_name:str="Colbert-Experiment",load_only:bool=False,nranks:int=1,nbits:int=2,DOC_MAXLEN:int=300,INDEX_BSIZE:int=256,KMEANS_ITER:int=8):
+
 
  self.checkpoint = checkpoint
+ self.index_name_or_path = index_name_or_path
+ self.experiment_name = experiment_name
+ self.nranks = nranks
+ self.nbits = nbits
+ self.DOC_MAXLEN = DOC_MAXLEN
+ self.INDEX_BSIZE = INDEX_BSIZE
+ self.KMEANS_ITER = KMEANS_ITER
+ self.passages = passages
+
+ if not load_only:
+ print(f"Building the index for experiment {self.experiment_name} with index name {self.index_name_or_path}")
+ self.build_index()
+
+ print(f"Loading the index for experiment {self.experiment_name} with index name {self.index_name_or_path}")
+ self.searcher = self.get_index()
 
+ def build_index(self):
 
- def build_index(self,passages:List[str],nranks:int=1,index_name_or_path:str = "Colbert-RM-",nbits:int=2,DOC_MAXLEN:int=300,INDEX_BSIZE:int=256,KMEANS_ITER:int=8,experiment_name:str="Colbert-Experiment"):
+ try:
+ import colbert
+ except ImportError:
+ print("Colbert not found. Please check your installation or install the module using pip install colbert-ai[faiss-gpu,torch].")
 
- with Run().context(RunConfig(nranks=nranks, experiment=experiment_name)): 
- config = ColBERTConfig(doc_maxlen=DOC_MAXLEN, nbits=nbits, kmeans_niters=KMEANS_ITER,index_bsize=INDEX_BSIZE)
+ from colbert import Indexer
+ from colbert.infra import Run, RunConfig, ColBERTConfig
+ with Run().context(RunConfig(nranks=self.nranks, experiment=self.experiment_name)): 
+ config = ColBERTConfig(doc_maxlen=self.DOC_MAXLEN, nbits=self.nbits, kmeans_niters=self.KMEANS_ITER,index_bsize=self.INDEX_BSIZE)
 
 
  indexer = Indexer(checkpoint=self.checkpoint, config=config)
- indexer.index(name=index_name_or_path, collection=passages, overwrite=True)
+ indexer.index(name=self.index_name_or_path, collection=self.passages, overwrite=True)
 
- def get_index(self,index_name_or_path:str = "Colbert-RM-",experiment_name:str="Colbert-Experiment",passages:List[str] = []):
- with Run().context(RunConfig(experiment=experiment_name)):
- searcher = Searcher(index=index_name_or_path, collection=passages)
- self.searcher = searcher
+ def get_index(self):
+ try:
+ import colbert
+ except ImportError:
+ print("Colbert not found. Please check your installation or install the module using pip install colbert-ai[faiss-gpu,torch].")
+
+ from colbert import Searcher
+ from colbert.infra import Run, RunConfig
+
+ with Run().context(RunConfig(experiment=self.experiment_name)):
+ searcher = Searcher(index=self.index_name_or_path, collection=self.passages)
  return searcher
 
- def get_docs(self,searcher:Searcher,query:str,k:int=7):
+ def __call__(self,query:str,k:int=7,**kwargs):
+ try:
+ import colbert
+ except ImportError:
+ print("Colbert not found. Please check your installation or install the module using pip install colbert-ai[faiss-gpu,torch].")
+ import torch
 
- results = searcher.search(
- query,
- #Number of passages to receive
- k=k)
- #Passing the filter function of relevant 
- # filter_fn=lambda pids: torch.tensor(
- # [pid for pid in pids if pid in relevant_ids],dtype=torch.int32).to(device))
+ if kwargs.get("filtered_pids"):
+ filtered_pids = kwargs.get("filtered_pids")
+ assert type(filtered_pids) == List[int], "The filtered pids should be a list of integers"
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+ results = self.searcher.search(
+ query,
+ #Number of passages to receive
+ k=k, 
+ #Passing the filter function of relevant 
+ filter_fn=lambda pids: torch.tensor(
+ [pid for pid in pids if pid in filtered_pids],dtype=torch.int32).to(device))
+
  return results

rm_test.py

@@ -158,8 +158,8 @@
  "The early bird catches the worm."
 ]
 
- col = dspy.ColBERTv2Local()
- col.build_index(passages=passages)
- searcher = col.get_index(passages=passages[:10])
- res = searcher.get_docs(searcher,query="Software",k=5)
- print(res)
+ col = dspy.ColBERTv2Local(passages=passages)
+
+ # searcher = col.get_index(passages=passages[:10])
+ # res = searcher.get_docs(searcher,query="Software",k=5)
+ # print(res)