update: integrate FigureAnnotatorFromPageImage into MedQAAssistant

medrag_multi_modal/assistant/figure_annotation.py

@@ -92,44 +92,48 @@ Here are some clues you need to follow:
         )
 
     @weave.op()
-    def predict(self, image_artifact_address: str):
+    def predict(self, page_idx: int, image_artifact_address: str):
         """
-        Predicts figure annotations for images in a given artifact directory.
+        Predicts figure annotations for a specific page in a document.
 
-        This function retrieves an artifact directory using the provided image artifact address.
-        It reads metadata from a JSONL file in the artifact directory and iterates over each item in the metadata.
-        For each item, it constructs the file path for the page image and checks for the presence of figure image files.
-        If figure image files are found, it reads and converts the page image, then uses the `annotate_figures` method
-        to extract figure annotations from the page image. The extracted annotations are then structured using the
-        `extract_structured_output` method and appended to the annotations list.
+        This function retrieves the artifact directory from the given image artifact address,
+        reads the metadata from the 'metadata.jsonl' file, and iterates through the metadata
+        to find the specified page index. If the page index matches, it reads the page image
+        and associated figure images, and then uses the `annotate_figures` method to extract
+        figure annotations from the page image. The extracted annotations are then structured
+        using the `extract_structured_output` method and returned as a dictionary.
 
         Args:
-            image_artifact_address (str): The address of the image artifact.
+            page_idx (int): The index of the page to annotate.
+            image_artifact_address (str): The address of the image artifact containing the page images.
 
         Returns:
-            list: A list of dictionaries containing page indices and their corresponding figure annotations.
+            dict: A dictionary containing the page index as the key and the extracted figure annotations
+                  as the value.
         """
         artifact_dir = get_wandb_artifact(image_artifact_address, "dataset")
         metadata = read_jsonl_file(os.path.join(artifact_dir, "metadata.jsonl"))
-        annotations = []
+        annotations = {}
         for item in track(metadata, description="Annotating images:"):
-            page_image_file = os.path.join(artifact_dir, f"page{item['page_idx']}.png")
-            figure_image_files = glob(
-                os.path.join(artifact_dir, f"page{item['page_idx']}_fig*.png")
-            )
-            if len(figure_image_files) > 0:
-                page_image = cv2.imread(page_image_file)
-                page_image = cv2.cvtColor(page_image, cv2.COLOR_BGR2RGB)
-                page_image = Image.fromarray(page_image)
-                figure_extracted_annotations = self.annotate_figures(
-                    page_image=page_image
+            if item["page_idx"] == page_idx:
+                page_image_file = os.path.join(
+                    artifact_dir, f"page{item['page_idx']}.png"
                 )
-                annotations.append(
-                    {
-                        "page_idx": item["page_idx"],
-                        "annotations": self.extract_structured_output(
-                            figure_extracted_annotations["annotations"]
-                        ).model_dump(),
-                    }
+                figure_image_files = glob(
+                    os.path.join(artifact_dir, f"page{item['page_idx']}_fig*.png")
                 )
+                if len(figure_image_files) > 0:
+                    page_image = cv2.imread(page_image_file)
+                    page_image = cv2.cvtColor(page_image, cv2.COLOR_BGR2RGB)
+                    page_image = Image.fromarray(page_image)
+                    figure_extracted_annotations = self.annotate_figures(
+                        page_image=page_image
+                    )
+                    figure_extracted_annotations = self.extract_structured_output(
+                        figure_extracted_annotations["annotations"]
+                    ).model_dump()
+                    annotations[item["page_idx"]] = figure_extracted_annotations[
+                        "annotations"
+                    ]
+                break
         return annotations

medrag_multi_modal/assistant/medqa_assistant.py

@@ -1,6 +1,9 @@
+from typing import Optional
+
 import weave
 
 from ..retrieval import SimilarityMetric
+from .figure_annotation import FigureAnnotatorFromPageImage
 from .llm_client import LLMClient
 
 
@@ -9,11 +12,12 @@ class MedQAAssistant(weave.Model):
 
     llm_client: LLMClient
     retriever: weave.Model
+    figure_annotator: FigureAnnotatorFromPageImage
     top_k_chunks: int = 2
     retrieval_similarity_metric: SimilarityMetric = SimilarityMetric.COSINE
 
     @weave.op()
-    def predict(self, query: str) -> str:
+    def predict(self, query: str, image_artifact_address: Optional[str] = None) -> str:
         retrieved_chunks = self.retriever.predict(
             query, top_k=self.top_k_chunks, metric=self.retrieval_similarity_metric
         )
@@ -23,13 +27,24 @@ class MedQAAssistant(weave.Model):
         for chunk in retrieved_chunks:
             retrieved_chunk_texts.append(chunk["text"])
             page_indices.add(int(chunk["page_idx"]))
-        page_numbers = ", ".join(map(str, page_indices))
+
+        figure_descriptions = []
+        if image_artifact_address is not None:
+            for page_idx in page_indices:
+                figure_annotations = self.figure_annotator.predict(
+                    page_idx=page_idx, image_artifact_address=image_artifact_address
+                )
+                figure_descriptions += [
+                    item["figure_description"] for item in figure_annotations[page_idx]
+                ]
 
         system_prompt = """
         You are an expert in medical science. You are given a query and a list of chunks from a medical document.
         """
         response = self.llm_client.predict(
-            system_prompt=system_prompt, user_prompt=[query, *retrieved_chunk_texts]
+            system_prompt=system_prompt,
+            user_prompt=[query, *retrieved_chunk_texts, *figure_descriptions],
         )
+        page_numbers = ", ".join([str(int(page_idx) + 1) for page_idx in page_indices])
         response += f"\n\n**Source:** {'Pages' if len(page_numbers) > 1 else 'Page'} {page_numbers} from Gray's Anatomy"
         return response