return metadata changes

dsp/primitives/search.py

@@ -14,7 +14,7 @@ def retrieve(query: str, k: int, **kwargs) -> list[str]:
  # 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]
+ # passages = [psg.long_text for psg in passages]
 
  if dsp.settings.reranker:
  passages_cs_scores = dsp.settings.reranker(query, passages)
@@ -55,17 +55,22 @@ def retrieveEnsemble(queries: list[str], k: int, by_prob: bool = True,**kwargs)
 
  if len(queries) == 1:
  return retrieve(queries[0], k)
-
- passages = {}
+ all_queries_passages = []
  for q in queries:
- for psg in dsp.settings.rm(q, k=k * 3,**kwargs):
+ passages = {}
+ retrieved_passages = dsp.settings.rm(q, k=k * 3,**kwargs)
+ # for idx,psg in enumerate(retrieved_passages):
+ # retrieved_passages[idx]["tracking_idx"] = idx
+ for idx,psg in enumerate(retrieved_passages):
  if by_prob:
- passages[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[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
+ passages[(idx,psg.long_text)] = passages.get(psg.long_text, 0.0) + psg.score
+ retrieved_passages[idx]["tracking_idx"] = idx
+ # passages = [(score, text) for text, score in passages.items()]
+ passages = sorted(passages.items(), key=lambda item: item[1])[:k]
+ # passages = sorted(passages, reverse=True)[: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]
+ all_queries_passages.append(passages)
+ return all_queries_passages

dspy/retrieve/retrieve.py

@@ -1,5 +1,5 @@
 import random
-from typing import List, Optional, Union
+from typing import List, Optional, Union, Dict, Any
 
 import dsp
 from dspy.predict.parameter import Parameter
@@ -37,6 +37,27 @@ def forward(self, query_or_queries: Union[str, List[str]], k: Optional[int] = No
  # 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)
+ if isinstance(passages[0],List):
+ pred_returns = []
+ for query_passages in passages:
+ passages_dict = {key:[] for key in list(query_passages[0].keys()) if key!="tracking_idx"}
+ for psg in query_passages:
+ for key,value in psg.items():
+ if key == "tracking_idx": continue
+ passages_dict[key].append(value)
+ if "long_text" in passages_dict:
+ passages_dict["passages"] = passages_dict.pop("long_text")
+ pred_returns.append(Prediction(**passages_dict)) 
+ 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)
 
 # TODO: Consider doing Prediction.from_completions with the individual sets of passages (per query) too.