Update app.py

app.py

@@ -2,8 +2,7 @@ import os
 import json
 import gradio as gr
 import pandas as pd
-import tempfile
-from typing import List
+from tempfile import NamedTemporaryFile
 
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_community.vectorstores import FAISS
@@ -11,31 +10,25 @@ from langchain_community.document_loaders import PyPDFLoader
 from langchain_core.output_parsers import StrOutputParser
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.llms import HuggingFaceHub
-from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_core.runnables import RunnableParallel, RunnablePassthrough
-from langchain_core.documents import Document
 
 huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
 
-def load_and_split_document(file: tempfile._TemporaryFileWrapper) -> List[Document]:
-    """Loads and splits the document into chunks."""
+def load_and_split_document(file):
+    """Loads and splits the document into pages."""
     loader = PyPDFLoader(file.name)
-    pages = loader.load()
-    
-    text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=1000,
-        chunk_overlap=200,
-        length_function=len,
-    )
-    
-    chunks = text_splitter.split_documents(pages)
-    return chunks
+    data = loader.load_and_split()
+    return data
 
 def get_embeddings():
     return HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
 
-def create_database(data: List[Document], embeddings):
-    db = FAISS.from_documents(data, embeddings)
+def create_or_update_database(data, embeddings):
+    if os.path.exists("faiss_database"):
+        db = FAISS.load_local("faiss_database", embeddings)
+        db.add_documents(data)
+    else:
+        db = FAISS.from_documents(data, embeddings)
     db.save_local("faiss_database")
 
 prompt = """
@@ -74,13 +67,19 @@ def response(database, model, question):
     ans = generate_chunked_response(model, formatted_prompt)
     return ans
 
-def update_vectors(file):
-    if file is None:
-        return "Please upload a PDF file."
-    data = load_and_split_document(file)
+def update_vectors(files):
+    if not files:
+        return "Please upload at least one PDF file."
+    
     embed = get_embeddings()
-    create_database(data, embed)
-    return f"Vector store updated successfully. Processed {len(data)} chunks."
+    total_chunks = 0
+    
+    for file in files:
+        data = load_and_split_document(file)
+        create_or_update_database(data, embed)
+        total_chunks += len(data)
+    
+    return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files."
 
 def ask_question(question):
     if not question:
@@ -98,7 +97,7 @@ def extract_db_to_excel():
     data = [{"page_content": doc.page_content, "metadata": json.dumps(doc.metadata)} for doc in documents]
     df = pd.DataFrame(data)
     
-    with tempfile.NamedTemporaryFile(delete=False, suffix='.xlsx') as tmp:
+    with NamedTemporaryFile(delete=False, suffix='.xlsx') as tmp:
         excel_path = tmp.name
         df.to_excel(excel_path, index=False)
     
@@ -109,7 +108,7 @@ with gr.Blocks() as demo:
     gr.Markdown("# Chat with your PDF documents")
     
     with gr.Row():
-        file_input = gr.File(label="Upload your PDF document", file_types=[".pdf"])
+        file_input = gr.File(label="Upload your PDF documents", file_types=[".pdf"], multiple=True)
         update_button = gr.Button("Update Vector Store")
     
     update_output = gr.Textbox(label="Update Status")