Merge pull request #39 from DL4DS/text_extraction

SCRUM-57 LLamaParser patch, DataLoader restructure and more

code/.chainlit/config.toml

@@ -23,7 +23,7 @@ allow_origins = ["*"]
 unsafe_allow_html = false
 
 # Process and display mathematical expressions. This can clash with "$" characters in messages.
-latex = false
+latex = true
 
 # Automatically tag threads with the current chat profile (if a chat profile is used)
 auto_tag_thread = true
@@ -85,31 +85,34 @@ custom_meta_image_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/f/
 # custom_build = "./public/build"
 
 [UI.theme]
-    default = "light"
+    default = "dark"
     #layout = "wide"
     #font_family = "Inter, sans-serif"
 # Override default MUI light theme. (Check theme.ts)
 [UI.theme.light]
-    background = "#FAFAFA"
-    paper = "#FFFFFF"
+    #background = "#FAFAFA"
+    #paper = "#FFFFFF"
 
     [UI.theme.light.primary]
-        main = "#b22222"  # Brighter shade of red
-        dark = "#8b0000"  # Darker shade of the brighter red
-        light = "#ff6347"  # Lighter shade of the brighter red
+        #main = "#F80061"
+        #dark = "#980039"
+        #light = "#FFE7EB"
     [UI.theme.light.text]
-        primary = "#212121"
-        secondary = "#616161"
+        #primary = "#212121"
+        #secondary = "#616161"
+
 # Override default MUI dark theme. (Check theme.ts)
 [UI.theme.dark]
-    background = "#1C1C1C" # Slightly lighter dark background color
-    paper = "#2A2A2A"      # Slightly lighter dark paper color
+    #background = "#FAFAFA"
+    #paper = "#FFFFFF"
 
     [UI.theme.dark.primary]
-        main = "#89CFF0"    # Primary color
-        dark = "#3700B3"    # Dark variant of primary color
-        light = "#CFBCFF"   # Lighter variant of primary color
-
+        #main = "#F80061"
+        #dark = "#980039"
+        #light = "#FFE7EB"
+    [UI.theme.dark.text]
+        #primary = "#EEEEEE"
+        #secondary = "#BDBDBD"
 
 [meta]
-generated_by = "1.1.302"
+generated_by = "1.1.304"

code/main.py

@@ -173,4 +173,6 @@ async def main(message):
     answer_with_sources, source_elements, sources_dict = get_sources(res, answer)
     processor._process(message.content, answer, sources_dict)
 
+    answer_with_sources = answer_with_sources.replace("$$", "$")
+
     await cl.Message(content=answer_with_sources, elements=source_elements).send()

code/modules/config/config.yml

@@ -34,7 +34,7 @@ llm_params:
   local_llm_params:
     model: 'tiny-llama'
     temperature: 0.7
-  pdf_reader: 'llama' # str [llama, pymupdf]
+  pdf_reader: 'llama' # str [llama, pymupdf, gpt]
 
 chat_logging:
   log_chat: False # bool

code/modules/config/constants.py

@@ -15,7 +15,9 @@ opening_message = f"Hey, What Can I Help You With?\n\nYou can me ask me question
 # Prompt Templates
 
 openai_prompt_template = """Use the following pieces of information to answer the user's question.
-You are an intelligent chatbot designed to help students with questions regarding the course. Render math equations in LaTeX format between $$ signs, and explain the parameters and variables in the equations.
+You are an intelligent chatbot designed to help students with questions regarding the course. 
+Render math equations in LaTeX format between $ or $$ signs, stick to the parameter and variable icons found in your context. 
+Be sure to explain the parameters and variables in the equations.
 If you don't know the answer, just say that you don't know. 
 
 Context: {context}
@@ -26,8 +28,9 @@ Helpful answer:
 """
 
 openai_prompt_template_with_history = """Use the following pieces of information to answer the user's question.
-You are an intelligent chatbot designed to help students with questions regarding the course. Render math equations in LaTeX format between $$ signs.
-
+You are an intelligent chatbot designed to help students with questions regarding the course. 
+Render math equations in LaTeX format between $ or $$ signs, stick to the parameter and variable icons found in your context. 
+Be sure to explain the parameters and variables in the equations.
 If you don't know the answer, just say that you don't know, don't try to make up an answer.
 
 Use the history to answer the question if you can.

code/modules/dataloader/data_loader.py

@@ -25,101 +25,19 @@ 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
+    from modules.dataloader.helpers import get_metadata, download_pdf_from_url
     from modules.config.constants import OPENAI_API_KEY, LLAMA_CLOUD_API_KEY
-
-
 except:
-    from dataloader.helpers import get_metadata
+    from dataloader.helpers import get_metadata, download_pdf_from_url
     from config.constants import OPENAI_API_KEY, LLAMA_CLOUD_API_KEY
 
 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):
-        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 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."
-        )
-
-    def parse(self, pdf_path):
-        pdf_name = os.path.basename(pdf_path)
-
-        documents = self.parser.load_data(pdf_path)
-        documents = [document.to_langchain_format() for document in documents]
-
-        os.remove(pdf_path) # cleanup, just in case
-        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):
@@ -186,18 +104,6 @@ class FileReader:
                 text += page.extract_text()
         return text
 
-    @staticmethod
-    def download_pdf_from_url(pdf_url):
-        response = requests.get(pdf_url)
-        if response.status_code == 200:
-            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
-                temp_file.write(response.content)
-                temp_file_path = temp_file.name
-            return temp_file_path
-        else:
-            self.logger.error(f"Failed to download PDF from URL: {pdf_url}")
-            return None
-
     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
@@ -383,22 +289,17 @@ 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
 
     def process_file(self, file_path, file_index, file_reader, addl_metadata):
         file_name = os.path.basename(file_path)
-        storage_dir = os.path.join(os.getcwd(), self.config["vectorstore"]["data_path"])
-        local_path = os.path.join(storage_dir, file_name)
-
-        if not os.path.exists(local_path):
-            local_path = FileReader.download_pdf_from_url(pdf_url=file_path)
 
         if file_name in self.document_data:
             return
 
-        file_type = file_name.split(".")[-1].lower()
-        self.logger.info(f"Reading file {file_index + 1}: {local_path}")
+        file_type = file_name.split(".")[-1]
 
         read_methods = {
             "pdf": file_reader.read_pdf,
@@ -412,9 +313,10 @@ class ChunkProcessor:
             return
 
         try:
-            documents = read_methods[file_type](local_path)
+            documents = read_methods[file_type](file_path)
+
             self.process_documents(
-                documents, local_path, file_type, "file", addl_metadata
+                documents, file_path, file_type, "file", addl_metadata
             )
         except Exception as e:
             self.logger.error(f"Error processing file {file_name}: {str(e)}")
@@ -500,10 +402,11 @@ if __name__ == "__main__":
     data_loader = DataLoader(config, logger=logger)
     document_chunks, document_names, documents, document_metadata = (
         data_loader.get_chunks(
-            uploaded_files,
-            ["https://dl4ds.github.io/sp2024/"],
+            ["https://dl4ds.github.io/sp2024/static_files/lectures/05_loss_functions_v2.pdf"],
+            [],
         )
     )
 
-    print(document_names)
+    print(document_names[:5])
     print(len(document_chunks))
+

code/modules/dataloader/helpers.py

@@ -1,7 +1,7 @@
 import requests
 from bs4 import BeautifulSoup
-from tqdm import tqdm
-
+from urllib.parse import urlparse
+import tempfile
 
 def get_urls_from_file(file_path: str):
     """
@@ -106,3 +106,23 @@ def get_metadata(lectures_url, schedule_url):
             continue
 
     return lecture_metadata
+
+
+def download_pdf_from_url(pdf_url):
+    """
+    Function to temporarily download a PDF file from a URL and return the local file path.
+
+    Args:
+        pdf_url (str): The URL of the PDF file to download.
+
+    Returns:
+        str: The local file path of the downloaded PDF file.
+    """
+    response = requests.get(pdf_url)
+    if response.status_code == 200:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
+            temp_file.write(response.content)
+            temp_file_path = temp_file.name
+        return temp_file_path
+    else:
+        return None

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)
+