app.py

import os
import json
import gradio as gr
import pandas as pd
from tempfile import NamedTemporaryFile

from langchain_core.prompts import ChatPromptTemplate
from langchain_community.vectorstores import FAISS
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_core.runnables import RunnableParallel, RunnablePassthrough

huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")

def load_and_split_document(file):
    """Loads and splits the document into pages."""
    loader = PyPDFLoader(file.name)
    data = loader.load_and_split()
    return data

def get_embeddings():
    return HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")

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 = """
Answer the question based only on the following context:
{context}
Question: {question}

Provide a concise and direct answer to the question:
"""

def get_model():
    return HuggingFaceHub(
        repo_id="mistralai/Mistral-7B-Instruct-v0.3",
        model_kwargs={"temperature": 0.5, "max_length": 512},
        huggingfacehub_api_token=huggingface_token
    )

def generate_chunked_response(model, prompt, max_tokens=500, max_chunks=5):
    full_response = ""
    for i in range(max_chunks):
        chunk = model(prompt + full_response, max_new_tokens=max_tokens)
        full_response += chunk
        if chunk.strip().endswith((".", "!", "?")):
            break
    return full_response.strip()

def response(database, model, question):
    prompt_val = ChatPromptTemplate.from_template(prompt)
    retriever = database.as_retriever()
    
    context = retriever.get_relevant_documents(question)
    context_str = "\n".join([doc.page_content for doc in context])
    
    formatted_prompt = prompt_val.format(context=context_str, question=question)
    
    ans = generate_chunked_response(model, formatted_prompt)
    return ans

def update_vectors(files):
    if not files:
        return "Please upload at least one PDF file."
    
    embed = get_embeddings()
    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:
        return "Please enter a question."
    embed = get_embeddings()
    database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
    model = get_model()
    return response(database, model, question)

def extract_db_to_excel():
    embed = get_embeddings()
    database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
    
    documents = database.docstore._dict.values()
    data = [{"page_content": doc.page_content, "metadata": json.dumps(doc.metadata)} for doc in documents]
    df = pd.DataFrame(data)
    
    with NamedTemporaryFile(delete=False, suffix='.xlsx') as tmp:
        excel_path = tmp.name
        df.to_excel(excel_path, index=False)
    
    return excel_path

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Chat with your PDF documents")
    
    with gr.Row():
        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")
    update_button.click(update_vectors, inputs=[file_input], outputs=update_output)
    
    with gr.Row():
        question_input = gr.Textbox(label="Ask a question about your documents")
        submit_button = gr.Button("Submit")
    
    answer_output = gr.Textbox(label="Answer")
    submit_button.click(ask_question, inputs=[question_input], outputs=answer_output)
    
    extract_button = gr.Button("Extract Database to Excel")
    excel_output = gr.File(label="Download Excel File")
    extract_button.click(extract_db_to_excel, inputs=[], outputs=excel_output)

if __name__ == "__main__":
    demo.launch()