Colbert PR fixes

dsp/modules/colbertv2.py

@@ -147,7 +147,7 @@ def forward(self,query:str,k:int=7,**kwargs):
  searcher_results = self.searcher.search(query, k=k)
  results = []
  for pid,rank,score in zip(*searcher_results):
- results.append(dotdict({'long_text':self.searcher.collection[pid],'pid':pid}))
+ results.append(dotdict({'long_text':self.searcher.collection[pid],'score':score,'pid':pid}))
  return results
 
 class ColBERTv2RerankerLocal:

dsp/primitives/search.py

@@ -1,17 +1,35 @@
-import warnings
+import logging
 from collections.abc import Iterable
 
 import numpy as np
 
 import dsp
 
+logger = logging.getLogger(__name__)
 
+# def retrieve(query: str, k: int, **kwargs) -> list[str]:
+# """Retrieves passages from the RM for the query and returns the top k passages."""
+# if not dsp.settings.rm:
+# raise AssertionError("No RM is loaded.")
+# passages = dsp.settings.rm(query, k=k, **kwargs)
+# if not isinstance(passages, Iterable):
+# # it's not an iterable yet; make it one.
+# # TODO: we should unify the type signatures of dspy.Retriever
+# passages = [passages]
+# passages = [psg.long_text for psg in passages]
+
+# if dsp.settings.reranker:
+# passages_cs_scores = dsp.settings.reranker(query, passages)
+# passages_cs_scores_sorted = np.argsort(passages_cs_scores)[::-1]
+# passages = [passages[idx] for idx in passages_cs_scores_sorted]
+
+
+# return passages
 def retrieve(query: str, k: int, **kwargs) -> list[str]:
  """Retrieves passages from the RM for the query and returns the top k passages."""
+
  if not dsp.settings.rm:
  raise AssertionError("No RM is loaded.")
- if not dsp.settings.reranker:
- warnings.warn("If you want to use the Reranker, please use dspy.RetrieveThenRerank",DeprecationWarning)
  passages = dsp.settings.rm(query, k=k, **kwargs)
  if not isinstance(passages, Iterable):
  # it's not an iterable yet; make it one.
@@ -21,50 +39,109 @@ def retrieve(query: str, k: int, **kwargs) -> list[str]:
  return passages
 
 
-def retrieveRerankEnsemble(queries: list[str], k: int,**kwargs) -> list[str]:
+# def retrieveRerankEnsemble(queries: list[str], k: int,**kwargs) -> list[str]:
+# if not (dsp.settings.rm and dsp.settings.reranker):
+# raise AssertionError("Both RM and Reranker are needed to retrieve & re-rank.")
+# queries = [q for q in queries if q]
+# passages = {}
+# for query in queries:
+# retrieved_passages = dsp.settings.rm(query, k=k*3,**kwargs)
+# passages_cs_scores = dsp.settings.reranker(query, [psg.long_text for psg in retrieved_passages])
+# for idx in np.argsort(passages_cs_scores)[::-1]:
+# psg = retrieved_passages[idx]
+# passages[psg.long_text] = passages.get(psg.long_text, []) + [
+# passages_cs_scores[idx],
+# ]
+
+
+# passages = [(np.average(score), text) for text, score in passages.items()]
+# return [text for _, text in sorted(passages, reverse=True)[:k]]
+def retrieveRerankEnsemble(queries: list[str], k: int, **kwargs) -> list[str]:
  if not (dsp.settings.rm and dsp.settings.reranker):
  raise AssertionError("Both RM and Reranker are needed to retrieve & re-rank.")
  queries = [q for q in queries if q]
  all_queries_passages = []
  for query in queries:
  passages = []
- retrieved_passages = dsp.settings.rm(query, k=k*3,**kwargs)
- passages_cs_scores = dsp.settings.reranker(query,passages=[psg["long_text"] for psg in retrieved_passages])
+ retrieved_passages = dsp.settings.rm(query, k=k * 3, **kwargs)
+ passages_cs_scores = dsp.settings.reranker(
+ query, passages=[psg["long_text"] for psg in retrieved_passages]
+ )
  for idx in np.argsort(passages_cs_scores)[::-1][:k]:
  curr_passage = retrieved_passages[idx]
- curr_passage['rerank_score'] = passages_cs_scores[idx]
+ curr_passage["rerank_score"] = passages_cs_scores[idx]
  passages.append(curr_passage)
  all_queries_passages.append(passages)
  if len(queries) == 1:
  return all_queries_passages[0]
  else:
  return all_queries_passages
 
-def retrieveEnsemble(queries: list[str], k: int, by_prob: bool = True,**kwargs) -> list[str]:
+
+# def retrieveEnsemble(queries: list[str], k: int, by_prob: bool = True,**kwargs) -> list[str]:
+# """Retrieves passages from the RM for each query in queries and returns the top k passages
+# based on the probability or score.
+# """
+# if not dsp.settings.rm:
+# raise AssertionError("No RM is loaded.")
+# if dsp.settings.reranker:
+# return retrieveRerankEnsemble(queries, k)
+
+# queries = [q for q in queries if q]
+
+# if len(queries) == 1:
+# return retrieve(queries[0], k, **kwargs)
+
+# passages = {}
+# for q in queries:
+# for psg in dsp.settings.rm(q, k=k * 3,**kwargs):
+# if by_prob:
+# passages[psg.long_text] = passages.get(psg.long_text, 0.0) + psg.prob
+# else:
+# passages[psg.long_text] = passages.get(psg.long_text, 0.0) + psg.score
+
+# passages = [(score, text) for text, score in passages.items()]
+# passages = sorted(passages, reverse=True)[:k]
+# passages = [text for _, text in passages]
+
+
+# return passages
+def retrieveEnsemble(
+ queries: list[str], k: int, by_prob: bool = True, **kwargs
+) -> list[str]:
  """Retrieves passages from the RM for each query in queries and returns the top k passages
  based on the probability or score.
  """
+
  if not dsp.settings.rm:
  raise AssertionError("No RM is loaded.")
  if not dsp.settings.reranker:
- warnings.warn("If you want to use the Reranker, please use dspy.RetrieveThenRerank. The reranking is ignored here.",DeprecationWarning)
-
+ logger.warn(
+ "DeprecationWarning: 'dspy.Retrieve' for reranking has been deprecated, please use dspy.RetrieveThenRerank. The reranking is ignored here. In the future this will raise an error."
+ )
+
  queries = [q for q in queries if q]
 
  if len(queries) == 1:
  return retrieve(queries[0], k)
  all_queries_passages = []
  for q in queries:
  passages = {}
- retrieved_passages =  dsp.settings.rm(q, k=k * 3,**kwargs)
- for idx,psg in enumerate(retrieved_passages):
+ retrieved_passages = dsp.settings.rm(q, k=k * 3, **kwargs)
+ for idx, psg in enumerate(retrieved_passages):
  if by_prob:
- passages[(idx,psg.long_text)] = passages.get(psg.long_text, 0.0) + psg.prob
+ passages[(idx, psg.long_text)] = (
+ passages.get(psg.long_text, 0.0) + psg.prob
+ )
  else:
- passages[(idx,psg.long_text)] = passages.get(psg.long_text, 0.0) + psg.score
+ passages[(idx, psg.long_text)] = (
+ passages.get(psg.long_text, 0.0) + psg.score
+ )
  retrieved_passages[idx]["tracking_idx"] = idx
  passages = sorted(passages.items(), key=lambda item: item[1])[:k]
  req_indices = [psg[0][0] for psg in passages]
- passages = [rp for rp in retrieved_passages if rp.get("tracking_idx") in req_indices]
+ passages = [
+ rp for rp in retrieved_passages if rp.get("tracking_idx") in req_indices
+ ]
  all_queries_passages.append(passages)
  return all_queries_passages

dspy/retrieve/retrieve.py

@@ -5,6 +5,14 @@
 from dspy.predict.parameter import Parameter
 from dspy.primitives.prediction import Prediction
 
+def single_query_passage(passages):
+ passages_dict = {key:[] for key in list(passages[0].keys())}
+ for docs in passages:
+ for key,value in docs.items():
+ passages_dict[key].append(value)
+ if "long_text" in passages_dict:
+ passages_dict["passages"] = passages_dict.pop("long_text")
+ return Prediction(**passages_dict)
 
 class Retrieve(Parameter):
  name = "Search"
@@ -30,6 +38,14 @@ def __call__(self, *args, **kwargs):
  return self.forward(*args, **kwargs)
 
  def forward(self, query_or_queries: Union[str, List[str]], k: Optional[int] = None,**kwargs) -> Union[Prediction,List[Prediction]]:
+ # queries = [query_or_queries] if isinstance(query_or_queries, str) else query_or_queries
+ # queries = [query.strip().split('\n')[0].strip() for query in queries]
+
+ # # print(queries)
+ # # TODO: Consider removing any quote-like markers that surround the query too.
+ # k = k if k is not None else self.k
+ # passages = dsp.retrieveEnsemble(queries, k=k,**kwargs)
+ # return Prediction(passages=passages)
  queries = [query_or_queries] if isinstance(query_or_queries, str) else query_or_queries
  queries = [query.strip().split('\n')[0].strip() for query in queries]
 
@@ -51,14 +67,7 @@ def forward(self, query_or_queries: Union[str, List[str]], k: Optional[int] = No
  return pred_returns
  elif isinstance(passages[0], Dict):
  #passages dict will contain {"long_text":long_text_list,"metadatas";metadatas_list...}
- passages_dict = {key:[] for key in list(passages[0].keys())}
-
- for psg in passages:
- for key,value in psg.items():
- passages_dict[key].append(value)
- if "long_text" in passages_dict:
- passages_dict["passages"] = passages_dict.pop("long_text")
- return Prediction(**passages_dict)
+ return single_query_passage(passages=passages)
 
 # TODO: Consider doing Prediction.from_completions with the individual sets of passages (per query) too.
 
@@ -106,11 +115,5 @@ def forward(self, query_or_queries: Union[str, List[str]], k: Optional[int] = No
  pred_returns.append(Prediction(**passages_dict)) 
  return pred_returns
  elif isinstance(passages[0], Dict):
- passages_dict = {key:[] for key in list(passages[0].keys())}
- for docs in passages:
- for key,value in docs.items():
- passages_dict[key].append(value)
- if "long_text" in passages_dict:
- passages_dict["passages"] = passages_dict.pop("long_text")
- return Prediction(**passages_dict)
+ return single_query_passage(passages=passages)