Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import fitz  # PyMuPDF
 from docx import Document
 from sentence_transformers import SentenceTransformer
 import faiss
@@ -30,6 +31,18 @@ def extract_text_from_docx(docx_path):
         print(f"Error extracting text from DOCX: {e}")
     return text
 
+# Function to extract text from a PDF document
+def extract_text_from_pdf(pdf_path):
+    text = ""
+    try:
+        pdf_document = fitz.open(pdf_path)
+        for page_num in range(pdf_document.page_count):
+            page = pdf_document.load_page(page_num)
+            text += page.get_text()
+    except Exception as e:
+        print(f"Error extracting text from PDF: {e}")
+    return text
+
 # Load or create FAISS index
 index_path = "faiss_index.pkl"
 document_texts_path = "document_texts.pkl"
@@ -64,13 +77,17 @@ def upload_files(files):
             file_path = file.name  # Get the file path from the NamedString object
             if file_path.endswith('.docx'):
                 text = extract_text_from_docx(file_path)
-
-                # Process the text and update FAISS index
-                sentences = text.split("\n")
-                sentences = [preprocess_text(sentence) for sentence in sentences if sentence.strip()]
-                embeddings = embedding_model.encode(sentences)
-                index.add(np.array(embeddings))
-                document_texts.extend(sentences)  # Store sentences for retrieval
+            elif file_path.endswith('.pdf'):
+                text = extract_text_from_pdf(file_path)
+            else:
+                continue
+
+            # Process the text and update FAISS index
+            sentences = text.split("\n")
+            sentences = [preprocess_text(sentence) for sentence in sentences if sentence.strip()]
+            embeddings = embedding_model.encode(sentences)
+            index.add(np.array(embeddings))
+            document_texts.extend(sentences)  # Store sentences for retrieval
 
         # Save the updated index and documents
         with open(index_path, "wb") as f:
@@ -110,14 +127,16 @@ def query_text(text):
         return f"Error querying text: {e}"
 
 def main():
-    upload_interface = gr.Interface(
+    # Gradio interface for uploading files
+    file_upload_interface = gr.Interface(
         fn=upload_files,
-        inputs=gr.File(file_count="multiple", label="Upload DOCX files"),
+        inputs=gr.File(file_count="multiple", label="Upload DOCX or PDF files"),
         outputs="text",
-        title="Upload DOCX Files",
-        description="Upload DOCX files to process and add to the FAISS index."
+        title="Upload Files",
+        description="Upload DOCX or PDF files to process and add to the FAISS index."
     )
 
+    # Gradio interface for querying text
     query_interface = gr.Interface(
         fn=query_text,
         inputs="text",
@@ -126,22 +145,9 @@ def main():
         description="Query the indexed text and get answers from the language model."
     )
 
-    demo = gr.TabbedInterface([upload_interface, query_interface], ["Upload Files", "Query Text"])
+    # Create a tabbed interface
+    demo = gr.TabbedInterface([file_upload_interface, query_interface], ["Upload Files", "Query Text"])
     demo.launch()
 
 if __name__ == "__main__":
     main()
-
-
-
-
-
-
-
-
-
-
-
-
-
-