Merge pull request #45 from DL4DS/text_extraction

PyMuPDF and HTML to Markdown fix + GPT4o mini PDF reader

code/modules/dataloader/data_loader.py

@@ -27,6 +27,7 @@ import tempfile
 import PyPDF2
 from modules.dataloader.pdf_readers.base import PDFReader
 from modules.dataloader.pdf_readers.llama import LlamaParser
+from modules.dataloader.pdf_readers.gpt import GPTParser
 
 try:
     from modules.dataloader.helpers import get_metadata, download_pdf_from_url
@@ -89,9 +90,12 @@ class FileReader:
         self.kind = kind
         if kind == "llama":
             self.pdf_reader = LlamaParser()
+        elif kind == "gpt":
+            self.pdf_reader = GPTParser()
         else:
             self.pdf_reader = PDFReader()
         self.web_reader = HTMLReader()
+        self.logger.info(f"Initialized FileReader with {kind} PDF reader and HTML reader")
 
 
     def extract_text_from_pdf(self, pdf_path):
@@ -105,11 +109,7 @@ class FileReader:
         return text
 
     def read_pdf(self, temp_file_path: str):
-        if self.kind == "llama":
-            documents = self.pdf_reader.parse(temp_file_path) # asyncio.run(self.pdf_reader.parse(temp_file_path)) if using async
-        else:
-            loader = self.pdf_reader.get_loader(temp_file_path)
-            documents = self.pdf_reader.get_documents(loader)
+        documents = self.pdf_reader.parse(temp_file_path)
         return documents
 
     def read_txt(self, temp_file_path: str):
@@ -134,8 +134,7 @@ class FileReader:
         return loader.load()
 
     def read_html(self, url: str):
-        loader = WebBaseLoader(url)
-        return loader.load()
+        return [Document(page_content=self.web_reader.read_html(url))]
 
     def read_tex_from_url(self, tex_url):
         response = requests.get(tex_url)
@@ -289,7 +288,6 @@ class ChunkProcessor:
                 )
                 self.document_chunks_full.extend(document_chunks)
 
-        print(f"Processed {file_path}. File_data: {file_data}")
         self.document_data[file_path] = file_data
         self.document_metadata[file_path] = file_metadata
 

code/modules/dataloader/pdf_readers/gpt.py

@@ -0,0 +1,81 @@
+import base64
+import os
+import requests
+
+from io import BytesIO
+from openai import OpenAI
+from pdf2image import convert_from_path
+from langchain.schema import Document
+
+
+class GPTParser:
+    """
+    This class uses OpenAI's GPT-4o mini model to parse PDFs and extract text, images and equations.
+    It is the most advanced parser in the system and is able to handle complex formats and layouts
+    """
+
+    def __init__(self):
+        self.client = OpenAI()
+        self.api_key = os.getenv("OPENAI_API_KEY")
+        self.prompt = """
+         The provided documents are images of 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 and convert everything to markdown format. Some of the equations may be complicated.
+         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. Separate each page with '---'.
+         Just respond with the markdown.
+         """
+
+    def parse(self, pdf_path):
+        images = convert_from_path(pdf_path)
+
+        encoded_images = [self.encode_image(image) for image in images]
+
+        chunks = [encoded_images[i:i + 5] for i in range(0, len(encoded_images), 5)]
+
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.api_key}"
+        }
+
+        output = ""
+        for chunk_num, chunk in enumerate(chunks):
+            content = [{"type": "image_url", "image_url": {
+                "url": f"data:image/jpeg;base64,{image}"}} for image in chunk]
+
+            content.insert(0, {"type": "text", "text": self.prompt})
+
+            payload = {
+                "model": "gpt-4o-mini",
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": content
+                    }
+                ],
+            }
+
+            response = requests.post(
+                "https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
+
+            resp = response.json()
+
+            chunk_output = resp['choices'][0]['message']['content'].replace("```", "").replace("markdown", "").replace("````", "")
+
+            output += chunk_output + "\n---\n"
+
+        output = output.split("\n---\n")
+        output = [doc for doc in output if doc.strip() != ""]
+
+        documents = [
+            Document(
+                page_content=page,
+                metadata={"source": pdf_path, "page": i}
+            ) for i, page in enumerate(output)
+        ]
+        return documents
+
+    def encode_image(self, image):
+        buffered = BytesIO()
+        image.save(buffered, format="JPEG")
+        return base64.b64encode(buffered.getvalue()).decode('utf-8')