Update app.py

app.py

@@ -3,7 +3,7 @@ import json
 import gradio as gr
 import pandas as pd
 from tempfile import NamedTemporaryFile
-
+from typing import List
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_community.vectorstores import FAISS
 from langchain_community.document_loaders import PyPDFLoader
@@ -11,119 +11,125 @@ 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
-
+from langchain_core.text_splitters import RecursiveCharacterTextSplitter
+from langchain_core.document import Document
 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 load_and_split_document_basic(file):
+   """Loads and splits the document into pages."""
+   loader = PyPDFLoader(file.name)
+   data = loader.load_and_split()
+   return data
+def load_and_split_document_recursive(file: NamedTemporaryFile) -> 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")
-
+   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, allow_dangerous_deserialization=True)
-        db.add_documents(data)
-    else:
-        db = FAISS.from_documents(data, embeddings)
-    db.save_local("faiss_database")
-
+   if os.path.exists("faiss_database"):
+       db = FAISS.load_local("faiss_database", embeddings, allow_dangerous_deserialization=True)
+       db.add_documents(data)
+   else:
+       db = FAISS.from_documents(data, embeddings)
+   db.save_local("faiss_database")
+def clear_cache():
+   if os.path.exists("faiss_database"):
+       os.remove("faiss_database")
+       return "Cache cleared successfully."
+   else:
+       return "No cache to clear."
 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.2, "max_length": 512},
-        huggingfacehub_api_token=huggingface_token
-    )
-
+def get_model(temperature, top_p, repetition_penalty):
+   return HuggingFaceHub(
+       repo_id="mistralai/Mistral-7B-Instruct-v0.3",
+       model_kwargs={
+           "temperature": temperature,
+           "top_p": top_p,
+           "repetition_penalty": repetition_penalty,
+           "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()
-
+   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)
-
+   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, use_recursive_splitter):
+   if not files:
+       return "Please upload at least one PDF file."
+   embed = get_embeddings()
+   total_chunks = 0
+   for file in files:
+       if use_recursive_splitter:
+           data = load_and_split_document_recursive(file)
+       else:
+           data = load_and_split_document_basic(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, temperature, top_p, repetition_penalty):
+   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(temperature, top_p, repetition_penalty)
+   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
-
+   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.Files(label="Upload your PDF documents", 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)
-
+   gr.Markdown("# Chat with your PDF documents")
+   with gr.Row():
+       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)
+   update_output = gr.Textbox(label="Update Status")
+   update_button.click(update_vectors, inputs=[file_input, use_recursive_splitter], outputs=update_output)
+   with gr.Row():
+       question_input = gr.Textbox(label="Ask a question about your documents")
+       temperature_slider = gr.Slider(label="Temperature", minimum=0.0, maximum=1.0, value=0.5, step=0.1)
+       top_p_slider = gr.Slider(label="Top P", minimum=0.0, maximum=1.0, value=0.9, step=0.1)
+       repetition_penalty_slider = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, value=1.0, step=0.1)
+       submit_button = gr.Button("Submit")
+   answer_output = gr.Textbox(label="Answer")
+   submit_button.click(ask_question, inputs=[question_input, temperature_slider, top_p_slider, repetition_penalty_slider], 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)
+   clear_button = gr.Button("Clear Cache")
+   clear_output = gr.Textbox(label="Cache Status")
+   clear_button.click(clear_cache, inputs=[], outputs=clear_output)
 if __name__ == "__main__":
-    demo.launch()
+   demo.launch()