PDF Reader Re-structure

code/modules/dataloader/data_loader.py

@@ -25,6 +25,8 @@ import html2text
 import bs4
 import tempfile
 import PyPDF2
+from modules.dataloader.pdf_readers.base import PDFReader
+from modules.dataloader.pdf_readers.llama import LlamaParser
 
 try:
     from modules.dataloader.helpers import get_metadata, download_pdf_from_url
@@ -36,103 +38,6 @@ except:
 logger = logging.getLogger(__name__)
 BASE_DIR = os.getcwd()
 
-class PDFReader:
-    def __init__(self):
-        pass
-
-    def get_loader(self, pdf_path):
-        loader = PyMuPDFLoader(pdf_path)
-        return loader
-
-    def get_documents(self, loader):
-        return loader.load()
-
-class LlamaParser:
-    def __init__(self):
-        print("Initializing LlamaParser")
-        self.GPT_API_KEY = OPENAI_API_KEY
-        self.LLAMA_CLOUD_API_KEY = LLAMA_CLOUD_API_KEY
-        self.parse_url = "https://api.cloud.llamaindex.ai/api/parsing/upload"
-        self.headers = {
-            'Accept': 'application/json',
-            'Authorization': f'Bearer {LLAMA_CLOUD_API_KEY}'
-        }
-        self.parser = LlamaParse(
-            api_key=LLAMA_CLOUD_API_KEY,
-            result_type="markdown",
-            verbose=True,
-            language="en",
-            gpt4o_mode=False,
-            # gpt4o_api_key=OPENAI_API_KEY,
-            parsing_instruction="The provided documents are PDFs of lecture slides of deep learning material. They contain LaTeX equations, images, and text. The goal is to extract the text, images and equations from the slides. The markdown should be clean and easy to read, and any math equation should be converted to LaTeX format, between $ signs. For images, if you can, give a description and a source."
-        )
-
-    def parse(self, pdf_path):
-        pdf_name = os.path.basename(pdf_path)
-
-        if not os.path.exists(pdf_path):
-            logger.warning(f"File {pdf_name} does not exist locally, installing temporarily...")
-            pdf_path = download_pdf_from_url(pdf_path)
-
-        documents = self.parser.load_data(pdf_path)
-        document = [document.to_langchain_format() for document in documents][0]
-
-        content = document.page_content
-        pages = content.split("\n---\n")
-        pages = [page.strip() for page in pages]
-
-        documents = [
-            Document(
-                page_content=page,
-                metadata={"source": pdf_path, "page": i}
-            ) for i, page in enumerate(pages)
-        ]
-
-        return documents
-
-    def make_request(self, pdf_url):
-        payload = {
-            "gpt4o_mode": "false",
-            "parsing_instruction": "The provided document is a PDF of lecture slides of deep learning material. They contain LaTeX equations, images, and text. The goal is to extract the text, images and equations from the slides and convert them to markdown format. The markdown should be clean and easy to read, and any math equation should be converted to LaTeX, between $$. For images, give a description and if you can, a source.",
-        }
-
-        files = [
-            ('file', ('file', requests.get(pdf_url).content, 'application/octet-stream'))
-        ]
-
-        response = requests.request(
-            "POST", self.parse_url, headers=self.headers, data=payload, files=files)
-
-        return response.json()['id'], response.json()['status']
-
-    async def get_result(self, job_id):
-        url = f"https://api.cloud.llamaindex.ai/api/parsing/job/{job_id}/result/markdown"
-
-        response = requests.request("GET", url, headers=self.headers, data={})
-
-        return response.json()['markdown']
-
-    async def _parse(self, pdf_path):
-        job_id, status = self.make_request(pdf_path)
-
-        while status != "SUCCESS":
-            url = f"https://api.cloud.llamaindex.ai/api/parsing/job/{job_id}"
-            response = requests.request("GET", url, headers=self.headers, data={})
-            status = response.json()["status"]
-
-        result = await self.get_result(job_id)
-
-        documents = [
-            Document(
-                page_content=result,
-                metadata={"source": pdf_path}
-            )
-        ]
-
-        return documents
-
-    async def _parse(self, pdf_path):
-        return await self._parse(pdf_path)
 
 class HTMLReader:
     def __init__(self):

code/modules/dataloader/pdf_readers/base.py

@@ -0,0 +1,14 @@
+from langchain_community.document_loaders import PyMuPDFLoader
+
+
+class PDFReader:
+    def __init__(self):
+        pass
+
+    def get_loader(self, pdf_path):
+        loader = PyMuPDFLoader(pdf_path)
+        return loader
+
+    def parse(self, pdf_path):
+        loader = self.get_loader(pdf_path)
+        return loader.load()

code/modules/dataloader/pdf_readers/llama.py

@@ -0,0 +1,92 @@
+import os
+import requests
+from llama_parse import LlamaParse
+from langchain.schema import Document
+from modules.config.constants import OPENAI_API_KEY, LLAMA_CLOUD_API_KEY
+from modules.dataloader.helpers import download_pdf_from_url
+
+
+
+class LlamaParser:
+    def __init__(self):
+        self.GPT_API_KEY = OPENAI_API_KEY
+        self.LLAMA_CLOUD_API_KEY = LLAMA_CLOUD_API_KEY
+        self.parse_url = "https://api.cloud.llamaindex.ai/api/parsing/upload"
+        self.headers = {
+            'Accept': 'application/json',
+            'Authorization': f'Bearer {LLAMA_CLOUD_API_KEY}'
+        }
+        self.parser = LlamaParse(
+            api_key=LLAMA_CLOUD_API_KEY,
+            result_type="markdown",
+            verbose=True,
+            language="en",
+            gpt4o_mode=False,
+            # gpt4o_api_key=OPENAI_API_KEY,
+            parsing_instruction="The provided documents are PDFs of lecture slides of deep learning material. They contain LaTeX equations, images, and text. The goal is to extract the text, images and equations from the slides. The markdown should be clean and easy to read, and any math equation should be converted to LaTeX format, between $ signs. For images, if you can, give a description and a source."
+        )
+
+    def parse(self, pdf_path):
+        if not os.path.exists(pdf_path):
+            pdf_path = download_pdf_from_url(pdf_path)
+
+        documents = self.parser.load_data(pdf_path)
+        document = [document.to_langchain_format() for document in documents][0]
+
+        content = document.page_content
+        pages = content.split("\n---\n")
+        pages = [page.strip() for page in pages]
+
+        documents = [
+            Document(
+                page_content=page,
+                metadata={"source": pdf_path, "page": i}
+            ) for i, page in enumerate(pages)
+        ]
+
+        return documents
+
+    def make_request(self, pdf_url):
+        payload = {
+            "gpt4o_mode": "false",
+            "parsing_instruction": "The provided document is a PDF of lecture slides of deep learning material. They contain LaTeX equations, images, and text. The goal is to extract the text, images and equations from the slides and convert them to markdown format. The markdown should be clean and easy to read, and any math equation should be converted to LaTeX, between $$. For images, give a description and if you can, a source.",
+        }
+
+        files = [
+            ('file', ('file', requests.get(pdf_url).content, 'application/octet-stream'))
+        ]
+
+        response = requests.request(
+            "POST", self.parse_url, headers=self.headers, data=payload, files=files)
+
+        return response.json()['id'], response.json()['status']
+
+    async def get_result(self, job_id):
+        url = f"https://api.cloud.llamaindex.ai/api/parsing/job/{job_id}/result/markdown"
+
+        response = requests.request("GET", url, headers=self.headers, data={})
+
+        return response.json()['markdown']
+
+    async def _parse(self, pdf_path):
+        job_id, status = self.make_request(pdf_path)
+
+        while status != "SUCCESS":
+            url = f"https://api.cloud.llamaindex.ai/api/parsing/job/{job_id}"
+            response = requests.request("GET", url, headers=self.headers, data={})
+            status = response.json()["status"]
+
+        result = await self.get_result(job_id)
+
+        documents = [
+            Document(
+                page_content=result,
+                metadata={"source": pdf_path}
+            )
+        ]
+
+        return documents
+
+    async def _parse(self, pdf_path):
+        return await self._parse(pdf_path)
+