Adding PDF

aimakerspace/text_utils.py

@@ -1,5 +1,6 @@
 import os
 from typing import List
+import PyPDF2
 
 
 class TextFileLoader:
@@ -62,6 +63,53 @@ class CharacterTextSplitter:
         return chunks
 
 
+class PDFLoader:
+    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.lower().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:
+            # Create PDF reader object
+            pdf_reader = PyPDF2.PdfReader(file)
+            
+            # Extract text from each page
+            text = ""
+            for page in pdf_reader.pages:
+                text += page.extract_text() + "\n"
+            
+            self.documents.append(text)
+
+    def load_directory(self):
+        for root, _, files in os.walk(self.path):
+            for file in files:
+                if file.lower().endswith('.pdf'):
+                    file_path = os.path.join(root, file)
+                    with open(file_path, 'rb') as f:
+                        pdf_reader = PyPDF2.PdfReader(f)
+                        
+                        # Extract text from each page
+                        text = ""
+                        for page in pdf_reader.pages:
+                            text += page.extract_text() + "\n"
+                        
+                        self.documents.append(text)
+
+    def load_documents(self):
+        self.load()
+        return self.documents
+
+
 if __name__ == "__main__":
     loader = TextFileLoader("data/KingLear.txt")
     loader.load()

pyproject.toml

@@ -9,5 +9,6 @@ dependencies = [
     "numpy>=2.2.2",
     "openai>=1.59.9",
     "pydantic==2.10.1",
+    "pypdf>=5.1.0",
     "websockets>=14.2",
-]
+]

uv.lock

@@ -10,6 +10,7 @@ dependencies = [
     { name = "numpy" },
     { name = "openai" },
     { name = "pydantic" },
+    { name = "pypdf" },
     { name = "websockets" },
 ]
 
@@ -19,6 +20,7 @@ requires-dist = [
     { name = "numpy", specifier = ">=2.2.2" },
     { name = "openai", specifier = ">=1.59.9" },
     { name = "pydantic", specifier = "==2.10.1" },
+    { name = "pypdf", specifier = ">=5.1.0" },
     { name = "websockets", specifier = ">=14.2" },
 ]
 
@@ -632,6 +634,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/61/ad/689f02752eeec26aed679477e80e632ef1b682313be70793d798c1d5fc8f/PyJWT-2.10.1-py3-none-any.whl", hash = "sha256:dcdd193e30abefd5debf142f9adfcdd2b58004e644f25406ffaebd50bd98dacb", size = 22997 },
 ]
 
+[[package]]
+name = "pypdf"
+version = "5.1.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/6b/9a/72d74f05f64895ebf1c7f6646cf7fe6dd124398c5c49240093f92d6f0fdd/pypdf-5.1.0.tar.gz", hash = "sha256:425a129abb1614183fd1aca6982f650b47f8026867c0ce7c4b9f281c443d2740", size = 5011381 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/04/fc/6f52588ac1cb4400a7804ef88d0d4e00cfe57a7ac6793ec3b00de5a8758b/pypdf-5.1.0-py3-none-any.whl", hash = "sha256:3bd4f503f4ebc58bae40d81e81a9176c400cbbac2ba2d877367595fb524dfdfc", size = 297976 },
+]
+
 [[package]]
 name = "python-dotenv"
 version = "1.0.1"