add: predict functions

medrag_multi_modal/retrieval/bm25s_retrieval.py

@@ -141,21 +141,6 @@ class BM25sRetriever(weave.Model):
         The results are returned as a list of dictionaries, each containing a chunk and
         its corresponding relevance score.
 
-        !!! example "Example Usage"
-            ```python
-            import weave
-            from dotenv import load_dotenv
-
-            from medrag_multi_modal.retrieval import BM25sRetriever
-
-            load_dotenv()
-            weave.init(project_name="ml-colabs/medrag-multi-modal")
-            retriever = BM25sRetriever.from_wandb_artifact(
-                index_artifact_address="ml-colabs/medrag-multi-modal/grays-anatomy-bm25s:v2"
-            )
-            retrieved_chunks = retriever.retrieve(query="What are Ribosomes?")
-            ```
-
         Args:
             query (str): The input query string to search for relevant chunks.
             top_k (int, optional): The number of top relevant chunks to retrieve. Defaults to 2.
@@ -177,3 +162,37 @@ class BM25sRetriever(weave.Model):
         ):
             retrieved_chunks.append({"chunk": chunk, "score": score})
         return retrieved_chunks
+
+    @weave.op()
+    def predict(self, query: str, top_k: int = 2):
+        """
+        Predicts the top-k most relevant chunks for a given query using the BM25 algorithm.
+
+        This function is a wrapper around the `retrieve` method. It takes an input query string,
+        tokenizes it using the BM25 tokenizer, and retrieves the top-k most relevant chunks from
+        the BM25 index. The results are returned as a list of dictionaries, each containing a chunk
+        and its corresponding relevance score.
+
+        !!! example "Example Usage"
+            ```python
+            import weave
+            from dotenv import load_dotenv
+
+            from medrag_multi_modal.retrieval import BM25sRetriever
+
+            load_dotenv()
+            weave.init(project_name="ml-colabs/medrag-multi-modal")
+            retriever = BM25sRetriever.from_wandb_artifact(
+                index_artifact_address="ml-colabs/medrag-multi-modal/grays-anatomy-bm25s:v2"
+            )
+            retrieved_chunks = retriever.predict(query="What are Ribosomes?")
+            ```
+
+        Args:
+            query (str): The input query string to search for relevant chunks.
+            top_k (int, optional): The number of top relevant chunks to retrieve. Defaults to 2.
+
+        Returns:
+            list: A list of dictionaries, each containing a retrieved chunk and its relevance score.
+        """
+        return self.retrieve(query, top_k)

medrag_multi_modal/retrieval/contriever_retrieval.py

@@ -170,22 +170,6 @@ class ContrieverRetriever(weave.Model):
         cosine similarity or Euclidean distance. The top-k chunks with the highest similarity scores
         are returned as a list of dictionaries, each containing a chunk and its corresponding score.
 
-        !!! example "Example Usage"
-            ```python
-            import weave
-            from dotenv import load_dotenv
-
-            from medrag_multi_modal.retrieval import ContrieverRetriever, SimilarityMetric
-
-            load_dotenv()
-            weave.init(project_name="ml-colabs/medrag-multi-modal")
-            retriever = ContrieverRetriever.from_wandb_artifact(
-                chunk_dataset_name="grays-anatomy-chunks:v0",
-                index_artifact_address="ml-colabs/medrag-multi-modal/grays-anatomy-contriever:v1",
-            )
-            scores = retriever.retrieve(query="What are Ribosomes?", metric=SimilarityMetric.COSINE)
-            ```
-
         Args:
             query (str): The input query string to search for relevant chunks.
             top_k (int, optional): The number of top relevant chunks to retrieve. Defaults to 2.
@@ -213,3 +197,44 @@ class ContrieverRetriever(weave.Model):
                 }
             )
         return retrieved_chunks
+
+    @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 function is a wrapper around the `retrieve` method. It takes an input query string,
+        retrieves the top-k most relevant chunks from the precomputed vector index based on the
+        specified similarity metric, and returns the results as a list of dictionaries, each containing
+        a chunk and its corresponding relevance score.
+
+        !!! example "Example Usage"
+            ```python
+            import weave
+            from dotenv import load_dotenv
+
+            from medrag_multi_modal.retrieval import ContrieverRetriever, SimilarityMetric
+
+            load_dotenv()
+            weave.init(project_name="ml-colabs/medrag-multi-modal")
+            retriever = ContrieverRetriever.from_wandb_artifact(
+                chunk_dataset_name="grays-anatomy-chunks:v0",
+                index_artifact_address="ml-colabs/medrag-multi-modal/grays-anatomy-contriever:v1",
+            )
+            scores = retriever.predict(query="What are Ribosomes?", metric=SimilarityMetric.COSINE)
+            ```
+
+        Args:
+            query (str): The input query string to search for relevant chunks.
+            top_k (int, optional): The number of top relevant chunks to retrieve. Defaults to 2.
+            metric (SimilarityMetric, optional): The similarity metric to use for scoring. Defaults to cosine similarity.
+
+        Returns:
+            list: A list of dictionaries, each containing a retrieved chunk and its relevance score.
+        """
+        return self.retrieve(query, top_k, metric)

medrag_multi_modal/retrieval/medcpt_retrieval.py

@@ -200,23 +200,6 @@ class MedCPTRetriever(weave.Model):
         cosine similarity or Euclidean distance. The top-k chunks with the highest similarity scores
         are returned as a list of dictionaries, each containing a chunk and its corresponding score.
 
-        !!! example "Example Usage"
-            ```python
-            import weave
-            from dotenv import load_dotenv
-
-            import wandb
-            from medrag_multi_modal.retrieval import MedCPTRetriever
-
-            load_dotenv()
-            weave.init(project_name="ml-colabs/medrag-multi-modal")
-            retriever = MedCPTRetriever.from_wandb_artifact(
-                chunk_dataset_name="grays-anatomy-chunks:v0",
-                index_artifact_address="ml-colabs/medrag-multi-modal/grays-anatomy-medcpt:v0",
-            )
-            retriever.retrieve(query="What are Ribosomes?")
-            ```
-
         Args:
             query (str): The input query string to search for relevant chunks.
             top_k (int, optional): The number of top relevant chunks to retrieve. Defaults to 2.
@@ -253,3 +236,44 @@ class MedCPTRetriever(weave.Model):
                 }
             )
         return retrieved_chunks
+
+    @weave.op()
+    def predict(
+        self,
+        query: str,
+        top_k: int = 2,
+        metric: SimilarityMetric = SimilarityMetric.COSINE,
+    ):
+        """
+        Predicts the most relevant chunks for a given query.
+
+        This function uses the `retrieve` method to find the top-k relevant chunks
+        from the dataset based on the input query. It allows specifying the number
+        of top relevant chunks to retrieve and the similarity metric to use for scoring.
+
+        !!! example "Example Usage"
+            ```python
+            import weave
+            from dotenv import load_dotenv
+
+            import wandb
+            from medrag_multi_modal.retrieval import MedCPTRetriever
+
+            load_dotenv()
+            weave.init(project_name="ml-colabs/medrag-multi-modal")
+            retriever = MedCPTRetriever.from_wandb_artifact(
+                chunk_dataset_name="grays-anatomy-chunks:v0",
+                index_artifact_address="ml-colabs/medrag-multi-modal/grays-anatomy-medcpt:v0",
+            )
+            retriever.predict(query="What are Ribosomes?")
+            ```
+
+        Args:
+            query (str): The input query string to search for relevant chunks.
+            top_k (int, optional): The number of top relevant chunks to retrieve. Defaults to 2.
+            metric (SimilarityMetric, optional): The similarity metric to use for scoring. Defaults to cosine similarity.
+
+        Returns:
+            list: A list of dictionaries, each containing a retrieved chunk and its relevance score.
+        """
+        return self.retrieve(query, top_k, metric)

medrag_multi_modal/retrieval/nv_embed_2.py

@@ -177,27 +177,6 @@ class NVEmbed2Retriever(weave.Model):
         cosine similarity or Euclidean distance. The top-k chunks with the highest similarity scores
         are returned as a list of dictionaries, each containing a chunk and its corresponding score.
 
-        !!! example "Example Usage"
-            ```python
-            import weave
-            from dotenv import load_dotenv
-
-            import wandb
-            from medrag_multi_modal.retrieval import NVEmbed2Retriever
-
-            load_dotenv()
-            weave.init(project_name="ml-colabs/medrag-multi-modal")
-            retriever = NVEmbed2Retriever(model_name="nvidia/NV-Embed-v2")
-            retriever.index(
-                chunk_dataset_name="grays-anatomy-chunks:v0",
-                index_name="grays-anatomy-nvembed2",
-            )
-            retriever = NVEmbed2Retriever.from_wandb_artifact(
-                chunk_dataset_name="grays-anatomy-chunks:v0",
-                index_artifact_address="ml-colabs/medrag-multi-modal/grays-anatomy-nvembed2:v0",
-            )
-            ```
-
         Args:
             query (list[str]): The input query strings to search for relevant chunks.
             top_k (int, optional): The number of top relevant chunks to retrieve.
@@ -273,6 +252,7 @@ class NVEmbed2Retriever(weave.Model):
             list: A list of dictionaries, each containing a retrieved chunk and its relevance score.
         """
         query = [
-            f"Instruct: Given a question, retrieve passages that answer the question\nQuery: {query}"
+            f"""Instruct: Given a question, retrieve passages that answer the question
+Query: {query}"""
         ]
         return self.retrieve(query, top_k, metric)