Merge pull request #2 from soumik12345/feat/text-loading

.gitignore

@@ -5,4 +5,5 @@ cursor_prompt.txt
 **pycache**
 .ruff_cache/
 test.py
-**.pdf
+**.pdf
+images/

docs/document_loader/load_text.md

@@ -1,3 +1,3 @@
 ## Load text from PDF files
 
-::: medrag_multi_modal.document_loader
+::: medrag_multi_modal.document_loader.load_text

docs/document_loader/load_text_image.md

@@ -0,0 +1,3 @@
+## Load text and images from PDF files
+
+::: medrag_multi_modal.document_loader.load_text_image

medrag_multi_modal/document_loader/__init__.py

@@ -1,3 +1,4 @@
-from .load_text import load_text_from_pdf
+from .load_text import TextLoader
+from .load_text_image import TextImageLoader
 
-__all__ = ["load_text_from_pdf"]
+__all__ = ["TextLoader", "TextImageLoader"]

medrag_multi_modal/document_loader/load_text.py

@@ -7,54 +7,33 @@ import PyPDF2
 import rich
 import weave
 from firerequests import FireRequests
-from pydantic import BaseModel
 
 
-class Page(BaseModel):
-    text: str
-    page_idx: int
-    document_name: str
-    file_path: str
-    file_url: str
-
-
-async def load_text_from_pdf(
-    url: str,
-    document_name: str,
-    document_file_path: str,
-    start_page: Optional[int] = None,
-    end_page: Optional[int] = None,
-    weave_dataset_name: Optional[str] = None,
-) -> list[Page]:
+class TextLoader:
     """
-    Asynchronously loads text from a PDF file specified by a URL or local file path,
-    processes the text into markdown format, and optionally publishes it to a Weave dataset.
+    A class for loading text from a PDF file, processing it into markdown, and optionally publishing it to a Weave dataset.
 
-    This function downloads a PDF from a given URL if it does not already exist locally,
-    reads the specified range of pages, converts each page's content to markdown, and
-    returns a list of Page objects containing the text and metadata. It uses PyPDF2 to read
-    the PDF and pymupdf4llm to convert pages to markdown. It processes pages concurrently using
-    `asyncio` for efficiency. If a weave_dataset_name is provided, the processed pages are published
-    to a Weave dataset.
+    This class handles the downloading of a PDF file from a given URL if it does not already exist locally.
+    It uses PyPDF2 to read the PDF and pymupdf4llm to convert pages to markdown. The processed pages are stored in a list
+    of Page objects, which can be optionally published to a Weave dataset.
 
-    !!! example "Example usage"
+    !!! example "Example Usage"
         ```python
         import asyncio
 
         import weave
 
-        from medrag_multi_modal.document_loader import load_text_from_pdf
+        from medrag_multi_modal.document_loader import TextLoader
 
         weave.init(project_name="ml-colabs/medrag-multi-modal")
         url = "https://archive.org/download/GraysAnatomy41E2015PDF/Grays%20Anatomy-41%20E%20%282015%29%20%5BPDF%5D.pdf"
+        loader = TextLoader(
+            url=url,
+            document_name="Gray's Anatomy",
+            document_file_path="grays_anatomy.pdf",
+        )
         asyncio.run(
-            load_text_from_pdf(
-                url=url,
-                document_name="Gray's Anatomy",
-                start_page=9,
-                end_page=15,
-                document_file_path="grays_anatomy.pdf",
-            )
+            loader.load_data(start_page=9, end_page=15, weave_dataset_name="grays-anatomy-text")
         )
         ```
 
@@ -62,61 +41,90 @@ async def load_text_from_pdf(
         url (str): The URL of the PDF file to download if not present locally.
         document_name (str): The name of the document for metadata purposes.
         document_file_path (str): The local file path where the PDF is stored or will be downloaded.
-        start_page (Optional[int]): The starting page index (0-based) to process. Defaults to the first page.
-        end_page (Optional[int]): The ending page index (0-based) to process. Defaults to the last page.
-        weave_dataset_name (Optional[str]): The name of the Weave dataset to publish the pages to, if provided.
-
-    Returns:
-        list[Page]: A list of Page objects, each containing the text and metadata for a processed page.
-
-    Raises:
-        ValueError: If the specified start_page or end_page is out of bounds of the document's page count.
     """
-    if not os.path.exists(document_file_path):
-        FireRequests().download(url, filename=document_file_path)
-    with open(document_file_path, "rb") as file:
-        pdf_reader = PyPDF2.PdfReader(file)
-        page_count = len(pdf_reader.pages)
-    print(f"Page count: {page_count}")
-    if start_page:
-        if start_page > page_count:
-            raise ValueError(
-                f"Start page {start_page} is greater than the total page count {page_count}"
-            )
-    else:
-        start_page = 0
-    if end_page:
-        if end_page > page_count:
-            raise ValueError(
-                f"End page {end_page} is greater than the total page count {page_count}"
-            )
-    else:
-        end_page = page_count - 1
-
-    pages: list[Page] = []
-    processed_pages_counter: int = 1
-    total_pages = end_page - start_page
 
-    async def process_page(page_idx):
-        nonlocal processed_pages_counter
-        text = pymupdf4llm.to_markdown(
-            doc=document_file_path, pages=[page_idx], show_progress=False
-        )
-        pages.append(
-            Page(
-                text=text,
-                page_idx=page_idx,
-                document_name=document_name,
-                file_path=document_file_path,
-                file_url=url,
+    def __init__(self, url: str, document_name: str, document_file_path: str):
+        self.url = url
+        self.document_name = document_name
+        self.document_file_path = document_file_path
+        if not os.path.exists(self.document_file_path):
+            FireRequests().download(url, filename=self.document_file_path)
+        with open(self.document_file_path, "rb") as file:
+            pdf_reader = PyPDF2.PdfReader(file)
+            self.page_count = len(pdf_reader.pages)
+
+    def get_page_indices(
+        self, start_page: Optional[int] = None, end_page: Optional[int] = None
+    ):
+        if start_page:
+            if start_page > self.page_count:
+                raise ValueError(
+                    f"Start page {start_page} is greater than the total page count {self.page_count}"
+                )
+        else:
+            start_page = 0
+        if end_page:
+            if end_page > self.page_count:
+                raise ValueError(
+                    f"End page {end_page} is greater than the total page count {self.page_count}"
+                )
+        else:
+            end_page = self.page_count - 1
+        return start_page, end_page
+
+    async def load_data(
+        self,
+        start_page: Optional[int] = None,
+        end_page: Optional[int] = None,
+        weave_dataset_name: Optional[str] = None,
+    ):
+        """
+        Asynchronously loads text from a PDF file specified by a URL or local file path,
+        processes the text into markdown format, and optionally publishes it to a Weave dataset.
+
+        This function downloads a PDF from a given URL if it does not already exist locally,
+        reads the specified range of pages, converts each page's content to markdown, and
+        returns a list of Page objects containing the text and metadata. It uses PyPDF2 to read
+        the PDF and pymupdf4llm to convert pages to markdown. It processes pages concurrently using
+        `asyncio` for efficiency. If a weave_dataset_name is provided, the processed pages are published
+        to a Weave dataset.
+
+        Args:
+            start_page (Optional[int]): The starting page index (0-based) to process. Defaults to the first page.
+            end_page (Optional[int]): The ending page index (0-based) to process. Defaults to the last page.
+            weave_dataset_name (Optional[str]): The name of the Weave dataset to publish the pages to, if provided.
+
+        Returns:
+            list[Page]: A list of Page objects, each containing the text and metadata for a processed page.
+
+        Raises:
+            ValueError: If the specified start_page or end_page is out of bounds of the document's page count.
+        """
+        start_page, end_page = self.get_page_indices(start_page, end_page)
+        pages = []
+        processed_pages_counter: int = 1
+        total_pages = end_page - start_page
+
+        async def process_page(page_idx):
+            nonlocal processed_pages_counter
+            text = pymupdf4llm.to_markdown(
+                doc=self.document_file_path, pages=[page_idx], show_progress=False
             )
-        )
-        rich.print(f"Processed pages {processed_pages_counter}/{total_pages}")
-        processed_pages_counter += 1
-
-    tasks = [process_page(page_idx) for page_idx in range(start_page, end_page)]
-    for task in asyncio.as_completed(tasks):
-        await task
-    if weave_dataset_name:
-        weave.publish(weave.Dataset(name=weave_dataset_name, rows=pages))
-    return pages
+            pages.append(
+                {
+                    "text": text,
+                    "page_idx": page_idx,
+                    "document_name": self.document_name,
+                    "file_path": self.document_file_path,
+                    "file_url": self.url,
+                }
+            )
+            rich.print(f"Processed pages {processed_pages_counter}/{total_pages}")
+            processed_pages_counter += 1
+
+        tasks = [process_page(page_idx) for page_idx in range(start_page, end_page)]
+        for task in asyncio.as_completed(tasks):
+            await task
+        if weave_dataset_name:
+            weave.publish(weave.Dataset(name=weave_dataset_name, rows=pages))
+        return pages

medrag_multi_modal/document_loader/load_text_image.py

@@ -0,0 +1,137 @@
+import asyncio
+import os
+from glob import glob
+from typing import Optional
+
+import pymupdf4llm
+import rich
+import weave
+from PIL import Image
+
+from medrag_multi_modal.document_loader.load_text import TextLoader
+
+
+class TextImageLoader(TextLoader):
+    """
+    A class for loading and processing text and images from a document.
+
+    The TextImageLoader class extends the TextLoader class to provide
+    functionality for extracting both text and images from a document
+    specified by a URL, document name, and file path. It processes the
+    document asynchronously, allowing for efficient handling of large
+    documents.
+
+    !!! example "Example Usage"
+        ```python
+        import asyncio
+
+        import weave
+
+        from medrag_multi_modal.document_loader import TextImageLoader
+
+        weave.init(project_name="ml-colabs/medrag-multi-modal")
+        url = "https://archive.org/download/GraysAnatomy41E2015PDF/Grays%20Anatomy-41%20E%20%282015%29%20%5BPDF%5D.pdf"
+        loader = TextImageLoader(
+            url=url,
+            document_name="Gray's Anatomy",
+            document_file_path="grays_anatomy.pdf",
+        )
+        asyncio.run(
+            loader.load_data(
+                start_page=20,
+                end_page=25,
+                weave_dataset_name="grays-anatomy-text",
+            )
+        )
+        ```
+
+    Args:
+        url (str): The URL of the document to be processed.
+        document_name (str): The name of the document.
+        document_file_path (str): The file path where the document is stored.
+    """
+
+    def __init__(self, url: str, document_name: str, document_file_path: str):
+        super().__init__(url, document_name, document_file_path)
+
+    async def load_data(
+        self,
+        start_page: Optional[int] = None,
+        end_page: Optional[int] = None,
+        weave_dataset_name: Optional[str] = None,
+        image_path: Optional[str] = "./images",
+        dpi: int = 300,
+    ):
+        """
+        Asynchronously loads and processes text and images from a specified range of pages
+        in a document. This function extracts text in markdown format and images in PNG
+        format from the document, storing them in a list of dictionaries, each representing
+        a page. Optionally, the processed data can be published to a Weave dataset.
+
+        The function first determines the page indices to process using the
+        `get_page_indices` method. It then defines an asynchronous inner function,
+        `process_page`, which handles the extraction of text and images for a single page.
+        The text is extracted using the `pymupdf4llm.to_markdown` function, and images are
+        retrieved from the specified image path. The processed data is appended to the
+        `pages` list.
+
+        The function creates a list of tasks for processing each page asynchronously and
+        awaits their completion. If a `weave_dataset_name` is provided, the processed data
+        is published to a Weave dataset. Finally, the function returns the list of processed
+        pages.
+
+        Args:
+            start_page (Optional[int]): The starting page index for processing. If None,
+                defaults to the first page of the document.
+            end_page (Optional[int]): The ending page index for processing. If None,
+                defaults to the last page of the document.
+            weave_dataset_name (Optional[str]): The name of the Weave dataset to publish
+                the processed data to. If None, the data is not published.
+            image_path (Optional[str]): The directory path where extracted images are
+                stored. Defaults to "./images".
+            dpi (int): The resolution in dots per inch for image extraction. Defaults to 300.
+
+        Returns:
+            List[Dict]: A list of dictionaries, each containing the extracted text, page
+            index, document name, file path, file URL, and a list of images for each page
+            processed.
+        """
+        start_page, end_page = self.get_page_indices(start_page, end_page)
+        pages = []
+        processed_pages_counter: int = 1
+        total_pages = end_page - start_page
+
+        async def process_page(page_idx):
+            nonlocal processed_pages_counter
+            text = pymupdf4llm.to_markdown(
+                doc=self.document_file_path,
+                pages=[page_idx],
+                show_progress=False,
+                write_images=True,
+                image_format="png",
+                dpi=dpi,
+                image_path=image_path,
+            )
+            image_paths = glob(
+                os.path.join(image_path, f"{self.document_file_path}-{page_idx}-*.png")
+            )
+            print(image_paths)
+            pages.append(
+                {
+                    "text": text,
+                    "page_idx": page_idx,
+                    "document_name": self.document_name,
+                    "file_path": self.document_file_path,
+                    "file_url": self.url,
+                    "images": [Image.open(image) for image in image_paths],
+                }
+            )
+            rich.print(f"Processed pages {processed_pages_counter}/{total_pages}")
+            processed_pages_counter += 1
+
+        tasks = [process_page(page_idx) for page_idx in range(start_page, end_page)]
+        for task in asyncio.as_completed(tasks):
+            await task
+        if weave_dataset_name:
+            weave.publish(weave.Dataset(name=weave_dataset_name, rows=pages))
+        return pages

mkdocs.yml

@@ -14,8 +14,6 @@ theme:
       toggle:
         icon: material/brightness-4
         name: Switch to light mode
-  logo: assets/logomark.svg
-  favicon: assets/logomark.svg
   features:
     - content.code.annotate
     - content.code.copy
@@ -63,5 +61,6 @@ nav:
   - Home: 'index.md'
   - Document Loader:
     - Text Loader: 'document_loader/load_text.md'
+    - Text and Image Loader: 'document_loader/load_text_image.md'
 
 repo_url: https://github.com/soumik12345/medrag-multi-modal