Update app.py

app.py

@@ -14,7 +14,7 @@ from langchain.prompts import PromptTemplate
 from langchain.chains import LLMChain
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_community.vectorstores import FAISS
-from langchain_community.document_loaders import PyPDFLoader
+from langchain_community.document_loaders import PyPDFLoader, PDFMinerLoader
 from langchain_core.output_parsers import StrOutputParser
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
@@ -25,8 +25,7 @@ from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
 from openpyxl import load_workbook
 from openpyxl.utils.dataframe import dataframe_to_rows
-
-
+import camelot
 
 huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
 
@@ -55,6 +54,86 @@ def load_and_split_document_recursive(file: NamedTemporaryFile) -> List[Document
     chunks = text_splitter.split_documents(pages)
     return chunks
 
+def load_and_split_document_basic(file: NamedTemporaryFile, parser: str) -> List[Document]:
+    """Loads and splits the document into pages."""
+    if parser == "PyPDF":
+        loader = PyPDFLoader(file.name)
+    elif parser == "PDFMiner":
+        loader = PDFMinerLoader(file.name)
+    elif parser == "Camelot":
+        return load_and_split_document_camelot(file)
+    else:
+        raise ValueError(f"Unknown parser: {parser}")
+    
+    return loader.load_and_split()
+
+def load_and_split_document_recursive(file: NamedTemporaryFile, parser: str) -> List[Document]:
+    """Loads and splits the document into chunks using recursive character text splitter."""
+    if parser == "PyPDF":
+        loader = PyPDFLoader(file.name)
+    elif parser == "PDFMiner":
+        loader = PDFMinerLoader(file.name)
+    elif parser == "Camelot":
+        return load_and_split_document_camelot(file)
+    else:
+        raise ValueError(f"Unknown parser: {parser}")
+    
+    pages = loader.load()
+    
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=1000,
+        chunk_overlap=200,
+        length_function=len,
+    )
+    
+    chunks = text_splitter.split_documents(pages)
+    return chunks
+
+def load_and_split_document_camelot(file: NamedTemporaryFile) -> List[Document]:
+    """Loads and splits the document using Camelot for tables and charts."""
+    tables = camelot.read_pdf(file.name, pages='all')
+    documents = []
+    
+    for i, table in enumerate(tables):
+        df = table.df
+        content = df.to_string(index=False)
+        documents.append(Document(page_content=content, metadata={"source": file.name, "table_number": i+1}))
+    
+    return documents
+
+def load_document(file: NamedTemporaryFile, parser: str, use_recursive_splitter: bool) -> List[Document]:
+    """Loads the document using the specified parser and splitting method."""
+    if parser == "Camelot":
+        return load_and_split_document_camelot(file)
+    elif use_recursive_splitter:
+        return load_and_split_document_recursive(file, parser)
+    else:
+        return load_and_split_document_basic(file, parser)
+
+def update_vectors(files, use_recursive_splitter, selected_parser):
+    if not files:
+        return "Please upload at least one PDF file."
+    
+    embed = get_embeddings()
+    total_chunks = 0
+    
+    all_data = []
+    for file in files:
+        data = load_document(file, selected_parser, use_recursive_splitter)
+        all_data.extend(data)
+        total_chunks += len(data)
+    
+    if os.path.exists("faiss_database"):
+        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+        database.add_documents(all_data)
+    else:
+        database = FAISS.from_documents(all_data, embed)
+    
+    database.save_local("faiss_database")
+    
+    splitting_method = "recursive splitting" if use_recursive_splitter else "page-by-page splitting"
+    return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {selected_parser} parser with {splitting_method}."
+
 def get_embeddings():
     return HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
 
@@ -682,7 +761,6 @@ def export_memory_db_to_excel():
     
     return excel_path
 
-# Gradio interface
 # Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Chat with your PDF documents and News")
@@ -691,9 +769,14 @@ with gr.Blocks() as demo:
         file_input = gr.Files(label="Upload your PDF documents", file_types=[".pdf"])
         update_button = gr.Button("Update Vector Store")
         use_recursive_splitter = gr.Checkbox(label="Use Recursive Text Splitter", value=False)
+        parser_dropdown = gr.Dropdown(
+            choices=["PyPDF", "PDFMiner", "Camelot"],
+            label="Select Parser",
+            value="PyPDF"
+        )
     
     update_output = gr.Textbox(label="Update Status")
-    update_button.click(update_vectors, inputs=[file_input, use_recursive_splitter], outputs=update_output)
+    update_button.click(update_vectors, inputs=[file_input, use_recursive_splitter, parser_dropdown], outputs=update_output)
     
     with gr.Row():
         with gr.Column(scale=2):