NVEmbed2Retriever.predict:v0
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import weave
from medrag_multi_modal.retrieval.common import SimilarityMetric
@weave.op()
def predict(
self,
query: str,
top_k: int = 2,
metric: SimilarityMetric = SimilarityMetric.COSINE,
):
query = [
f"Instruct: Given a question, retrieve passages that answer the question\nQuery: {query}"
]
return self.retrieve(query, top_k, metric)