DenseRetrieverWithReranker.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
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
retrievals = self.retriever.predict(query, top_k)
reranked = self.reranker.predict(query, retrievals, top_n)
return reranked