Adds pdf reader

aimakerspace/text_utils.py

@@ -1,5 +1,6 @@
 import os
 from typing import List
+from PyPDF2 import PdfReader
 
 
 class TextFileLoader:
@@ -35,6 +36,44 @@ class TextFileLoader:
         self.load()
         return self.documents
 
+class PDFFileLoader:
+    def __init__(self, path: str):
+        self.documents = []
+        self.path = path
+
+    def load(self):
+        if os.path.isdir(self.path):
+            self.load_directory()
+        elif os.path.isfile(self.path) and self.path.endswith(".pdf"):
+            self.load_file()
+        else:
+            raise ValueError(
+                "Provided path is neither a valid directory nor a .pdf file."
+            )
+
+    def load_file(self):
+        with open(self.path, "rb") as file:
+            pdf_reader = PdfReader(file)
+            text = ""
+            for page in pdf_reader.pages:
+                text += page.extract_text()
+            self.documents.append(text)
+
+    def load_directory(self):
+        for root, _, files in os.walk(self.path):
+            for file in files:
+                if file.endswith(".pdf"):
+                    file_path = os.path.join(root, file)
+                    with open(file_path, "rb") as f:
+                        pdf_reader = PdfReader(f)
+                        text = ""
+                        for page in pdf_reader.pages:
+                            text += page.extract_text()
+                        self.documents.append(text)
+
+    def load_documents(self):
+        self.load()
+        return self.documents
 
 class CharacterTextSplitter:
     def __init__(

app.py

@@ -1,7 +1,7 @@
 import os
 from typing import List
 from chainlit.types import AskFileResponse
-from aimakerspace.text_utils import CharacterTextSplitter, TextFileLoader
+from aimakerspace.text_utils import CharacterTextSplitter, TextFileLoader, PDFFileLoader
 from aimakerspace.openai_utils.prompts import (
     UserRolePrompt,
     SystemRolePrompt,
@@ -64,6 +64,19 @@ def process_text_file(file: AskFileResponse):
     texts = text_splitter.split_texts(documents)
     return texts
 
+def process_pdf_file(file: AskFileResponse):
+    import tempfile
+
+    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".pdf") as temp_file:
+        temp_file_path = temp_file.name
+
+    with open(temp_file_path, "wb") as f:
+        f.write(file.content)
+
+    text_loader = PDFFileLoader(temp_file_path)
+    documents = text_loader.load_documents()
+    texts = text_splitter.split_texts(documents)
+    return texts
 
 @cl.on_chat_start
 async def on_chat_start():
@@ -73,7 +86,7 @@ async def on_chat_start():
     while files == None:
         files = await cl.AskFileMessage(
             content="Please upload a Text File file to begin!",
-            accept=["text/plain"],
+            accept=["text/plain", "application/pdf"],
             max_size_mb=2,
             timeout=180,
         ).send()
@@ -86,7 +99,10 @@ async def on_chat_start():
     await msg.send()
 
     # load the file
-    texts = process_text_file(file)
+    if file.path.endswith(".pdf"):
+        texts = process_text_file(file)
+    else:
+        texts = process_pdf_file(file)
 
     print(f"Processing {len(texts)} text chunks")