app.py

import os
import json
import gradio as gr
import pandas as pd
import tempfile
from typing import List

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_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."""
    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

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)
    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(file):
    if file is None:
        return "Please upload a PDF file."
    data = load_and_split_document(file)
    embed = get_embeddings()
    create_database(data, embed)
    return f"Vector store updated successfully. Processed {len(data)} chunks."

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 tempfile.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 document", file_types=[".pdf"])
        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()