refactor: TextLoader

medrag_multi_modal/document_loader/__init__.py

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

medrag_multi_modal/document_loader/load_text.py

@@ -18,43 +18,31 @@ class Page(BaseModel):
     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.
-
-    !!! example "Example usage"
+    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"
         ```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 +50,91 @@ 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
+
+    @weave.op()
+    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: 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=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(
+                Page(
+                    text=text,
+                    page_idx=page_idx,
+                    document_name=self.document_name,
+                    file_path=self.document_file_path,
+                    file_url=self.url,
+                ).model_dump()
+            )
+            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