NVEmbed2Retriever.predict:v1
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
import weave
@weave.op()
def predict(
self,
query: str,
top_k: int = 2,
metric: SimilarityMetric = SimilarityMetric.COSINE,
):
"""
Predicts the top-k most relevant chunks for a given query using the specified similarity metric.
This method formats the input query string by prepending an instruction prompt and then calls the
`retrieve` method to get the most relevant chunks. The similarity metric can be either cosine similarity
or Euclidean distance. The top-k chunks with the highest similarity scores are returned.
!!! example "Example Usage"
```python
import weave
from medrag_multi_modal.retrieval.text_retrieval import NVEmbed2Retriever
weave.init(project_name="ml-colabs/medrag-multi-modal")
retriever = NVEmbed2Retriever.from_index(
index_repo_id="ashwiniai/medrag-text-corpus-chunks-nv-embed-2",