colbert integration

dsp/modules/__init__.py

@@ -4,7 +4,7 @@
 from .cache_utils import *
 from .clarifai import *
 from .cohere import *
-from .colbertv2 import ColBERTv2
+from .colbertv2 import ColBERTv2, ColBERTv2Local
 from .databricks import *
 from .google import *
 from .gpt3 import *

dsp/modules/colbertv2.py

@@ -1,10 +1,14 @@
 import functools
-from typing import Any, Optional, Union
+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
 
 # TODO: Ideally, this takes the name of the index and looks up its port.
 
@@ -74,3 +78,36 @@ def colbertv2_post_request_v2_wrapped(*args, **kwargs):
 
 
 colbertv2_post_request = colbertv2_post_request_v2_wrapped
+os.environ['COLBERT_LOAD_TORCH_EXTENSION_VERBOSE'] = "True"
+
+class ColBERTv2Local:
+ def __init__(self,checkpoint:str='colbert-ir/colbertv2.0'):
+
+ self.checkpoint = checkpoint
+
+
+ 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"):
+
+ 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)
+
+
+ indexer = Indexer(checkpoint=self.checkpoint, config=config)
+ indexer.index(name=index_name_or_path, collection=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
+ return searcher
+
+ def get_docs(self,searcher:Searcher,query:str,k:int=7):
+
+ 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))
+ return results

dspy/__init__.py

@@ -18,6 +18,7 @@
 Databricks = dsp.Databricks
 Cohere = dsp.Cohere
 ColBERTv2 = dsp.ColBERTv2
+ColBERTv2Local = dsp.ColBERTv2Local
 Pyserini = dsp.PyseriniRetriever
 Clarifai = dsp.ClarifaiLLM
 Google = dsp.Google

rm_test.py

@@ -25,18 +25,141 @@
 # print(retriever(["Software Internet"],by_prob=False,where={"table_name":"capexIndia"}))
 # print("-"*100)
 # print(retriever(["Software Internet","Packaging"],by_prob=False,where={"table_name":"capexIndia"}))
-import dspy
+# import dspy
+
+# colbertv2_wiki17_abstracts = dspy.ColBERTv2(url='http:https://20.102.90.50:2017/wiki17_abstracts')
+# dspy.settings.configure(rm=colbertv2_wiki17_abstracts,reranker=colbertv2_wiki17_abstracts)
 
-colbertv2_wiki17_abstracts = dspy.ColBERTv2(url='http:https://20.102.90.50:2017/wiki17_abstracts')
-dspy.settings.configure(rm=colbertv2_wiki17_abstracts,reranker=colbertv2_wiki17_abstracts)
+# #Define Retrieve Module
+# retriever = dspy.RetrieveThenRerank(k=3)
 
-#Define Retrieve Module
-retriever = dspy.RetrieveThenRerank(k=3)
+# query='When was the first FIFA World Cup held?'
 
-query='When was the first FIFA World Cup held?'
+# # Call the retriever on a particular query.
+# topK_passages = retriever([query])
 
-# Call the retriever on a particular query.
-topK_passages = retriever([query])
+# for idx, passage in enumerate(topK_passages):
+# print(f'{idx+1}]', passage, '\n')
+
+import os
+import dspy
+os.environ['COLBERT_LOAD_TORCH_EXTENSION_VERBOSE'] = "True"
+if __name__ == "__main__":
+ passages = [
+ "The quick brown fox jumps over the lazy dog.",
+ "She sells seashells by the seashore.",
+ "I am the master of my fate, I am the captain of my soul.",
+ "To be or not to be, that is the question.",
+ "All's fair in love and war.",
+ "A journey of a thousand miles begins with a single step.",
+ "Two wrongs don't make a right.",
+ "The pen is mightier than the sword.",
+ "Actions speak louder than words.",
+ "Beauty is in the eye of the beholder.",
+ "Practice makes perfect.",
+ "Where there's a will, there's a way.",
+ "When in Rome, do as the Romans do.",
+ "The early bird catches the worm.",
+ "You can't judge a book by its cover.",
+ "A picture is worth a thousand words.",
+ "Honesty is the best policy.",
+ "Don't count your chickens before they hatch.",
+ "Every cloud has a silver lining.",
+ "If at first you don't succeed, try, try again.",
+ "Look before you leap.",
+ "Rome wasn't built in a day.",
+ "The grass is always greener on the other side.",
+ "Absence makes the heart grow fonder.",
+ "Actions speak louder than words.",
+ "Ask and you shall receive.",
+ "Better late than never.",
+ "Don't bite the hand that feeds you.",
+ "Don't put all your eggs in one basket.",
+ "Easy come, easy go.",
+ "Every dog has its day.",
+ "Good things come to those who wait.",
+ "It's a piece of cake.",
+ "It's raining cats and dogs.",
+ "Kill two birds with one stone.",
+ "Let sleeping dogs lie.",
+ "Like father, like son.",
+ "Make hay while the sun shines.",
+ "Necessity is the mother of invention.",
+ "Out of sight, out of mind.",
+ "Patience is a virtue.",
+ "Practice what you preach.",
+ "The best things in life are free.",
+ "The squeaky wheel gets the grease.",
+ "There's no place like home.",
+ "Too many cooks spoil the broth.",
+ "When the going gets tough, the tough get going.",
+ "You reap what you sow.",
+ "A watched pot never boils.",
+ "Actions speak louder than words.",
+ "An apple a day keeps the doctor away.",
+ "Beggars can't be choosers.",
+ "Curiosity killed the cat.",
+ "Don't cry over spilled milk.",
+ "Don't put off until tomorrow what you can do today.",
+ "Every cloud has a silver lining.",
+ "Fortune favors the bold.",
+ "If the shoe fits, wear it.",
+ "It takes two to tango.",
+ "Keep your friends close and your enemies closer.",
+ "Let bygones be bygones.",
+ "No pain, no gain.",
+ "Once bitten, twice shy.",
+ "Practice makes perfect.",
+ "The apple doesn't fall far from the tree.",
+ "The early bird catches the worm.",
+ "The grass is always greener on the other side.",
+ "The more, the merrier.",
+ "There's no such thing as a free lunch.",
+ "To kill two birds with one stone.",
+ "When in Rome, do as the Romans do.",
+ "You can't have your cake and eat it too.",
+ "You can't make an omelet without breaking eggs.",
+ "A friend in need is a friend indeed.",
+ "A penny saved is a penny earned.",
+ "Actions speak louder than words.",
+ "Beauty is in the eye of the beholder.",
+ "Better late than never.",
+ "Don't count your chickens before they hatch.",
+ "Don't put all your eggs in one basket.",
+ "Every cloud has a silver lining.",
+ "If at first you don't succeed, try, try again.",
+ "If you can't beat them, join them.",
+ "Necessity is the mother of invention.",
+ "One man's trash is another man's treasure.",
+ "Practice makes perfect.",
+ "The early bird catches the worm.",
+ "The grass is always greener on the other side.",
+ "There's no place like home.",
+ "Too many cooks spoil the broth.",
+ "When in Rome, do as the Romans do.",
+ "You can't judge a book by its cover.",
+ "You reap what you sow.",
+ "A bird in the hand is worth two in the bush.",
+ "A penny for your thoughts.",
+ "Actions speak louder than words.",
+ "All good things must come to an end.",
+ "Beauty is only skin deep.",
+ "Don't bite the hand that feeds you.",
+ "Don't put off until tomorrow what you can do today.",
+ "Every dog has its day.",
+ "Fortune favors the bold.",
+ "If you want something done right, do it yourself.",
+ "It's better to be safe than sorry.",
+ "Make hay while the sun shines.",
+ "Necessity is the mother of invention.",
+ "Out of sight, out of mind.",
+ "Practice what you preach.",
+ "The best things in life are free.",
+ "The early bird catches the worm."
+]
 
-for idx, passage in enumerate(topK_passages):
- print(f'{idx+1}]', passage, '\n')
+ 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)