HybridRetrieverReranker.predict:v0
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
import weave
@weave.op()
def predict(self, query: str, top_k: int = None, top_n: int = None):
"""
Predicts the top-n results for the given query after re-ranking.
Args:
query (str): The search query.
top_k (int, optional): The number of top results to retrieve before re-ranking. Default is None.
top_n (int, optional): The number of top results to return after re-ranking. Default is None.
Returns:
list: A list of dictionaries containing the source, text, and score of the top-n results.
"""
if top_k and not top_n:
top_n = top_k
top_k = top_k * 2
elif top_n and not top_k:
top_k = top_n * 2
else:
top_k = 10
top_n = 5
sparse_retrievals = self.sparse_retriever.predict(query, top_k)
dense_retrievals = self.dense_retriever.predict(query, top_k)
fused = self.fusion_ranker.predict(sparse_retrievals, dense_retrievals)
reranked = self.ranker.predict(query, fused, top_n)
return reranked