Update app.py

app.py

@@ -1,29 +1,35 @@
 import gradio as gr
 import os
 import torch
-from langchain_community.vectorstores import FAISS
+from langchain_community.vectorstores import FAISS, Chroma
 from langchain_community.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.vectorstores import Chroma
 from langchain.chains import ConversationalRetrievalChain
-from langchain_community.embeddings import HuggingFaceEmbeddings 
-from langchain_community.llms import HuggingFacePipeline
-from langchain.chains import ConversationChain
-from langchain.memory import ConversationBufferMemory
+from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.llms import HuggingFaceEndpoint
+from langchain.memory import ConversationBufferMemory
+from langchain.retrievers import BM25Retriever, EnsembleRetriever
+from langchain.chains.query_constructor.base import AttributeInfo
+from langchain.chains import create_query_chain
+from langchain.retrievers.self_query.base import SelfQueryRetriever
+from langchain.chains.query_constructor.schema import FieldInfo
+from langchain.retrievers.multi_query import MultiQueryRetriever
 
 api_token = os.getenv("FirstToken")
 
 # Available LLM models
 list_llm = [
-    "meta-llama/Meta-Llama-3-8B-Instruct", 
+    "meta-llama/Meta-Llama-3-8B-Instruct",
     "mistralai/Mistral-7B-Instruct-v0.2",
     "deepseek-ai/deepseek-llm-7b-chat"
-]  
+]
 list_llm_simple = [os.path.basename(llm) for llm in list_llm]
 
+# -----------------------------------------------------------------------------
+# Document Loading and Splitting
+# -----------------------------------------------------------------------------
 def load_doc(list_file_path):
-    """Load and split PDF documents into chunks"""
+    """Load and split PDF documents into chunks."""
     loaders = [PyPDFLoader(x) for x in list_file_path]
     pages = []
     for loader in loaders:
@@ -35,14 +41,92 @@ def load_doc(list_file_path):
     doc_splits = text_splitter.split_documents(pages)
     return doc_splits
 
-def create_db(splits):
-    """Create vector database from document splits"""
+# -----------------------------------------------------------------------------
+# Vector Database Creation (ChromaDB and FAISS)
+# -----------------------------------------------------------------------------
+def create_chromadb(splits, persist_directory="chroma_db"):
+    """Create ChromaDB vector database from document splits."""
+    embeddings = HuggingFaceEmbeddings()
+    chromadb = Chroma.from_documents(
+        documents=splits,
+        embedding=embeddings,
+        persist_directory=persist_directory
+    )
+    chromadb.persist()  # Ensure data is written to disk
+    return chromadb
+
+def create_faissdb(splits):
+    """Create FAISS vector database from document splits."""
     embeddings = HuggingFaceEmbeddings()
-    vectordb = FAISS.from_documents(splits, embeddings)
-    return vectordb
+    faissdb = FAISS.from_documents(splits, embeddings)
+    return faissdb
+
+# -----------------------------------------------------------------------------
+# BM25 Retriever
+# -----------------------------------------------------------------------------
+def create_bm25_retriever(splits):
+    """Create BM25 retriever from document splits."""
+    bm25_retriever = BM25Retriever.from_documents(splits)
+    bm25_retriever.k = 3  # Number of documents to retrieve
+    return bm25_retriever
+
+# -----------------------------------------------------------------------------
+# MultiQueryRetriever
+# -----------------------------------------------------------------------------
+def create_multi_query_retriever(llm, vector_db, num_queries=3):
+    """
+    Create a MultiQueryRetriever.
+
+    Args:
+        llm: The language model to use for query generation.
+        vector_db: The vector database to retrieve from.
+        num_queries: The number of diverse queries to generate.
+
+    Returns:
+        A MultiQueryRetriever instance.
+    """
+    retriever = MultiQueryRetriever.from_llm(
+        llm=llm, retriever=vector_db.as_retriever(),
+        output_key="answer",
+        memory_key="chat_history",
+        return_messages=True,
+        verbose=False
+    )
+    return retriever
+
+# -----------------------------------------------------------------------------
+# Ensemble Retriever (Combine VectorDB and BM25)
+# -----------------------------------------------------------------------------
+def create_ensemble_retriever(vector_db, bm25_retriever):
+    """Create an ensemble retriever combining ChromaDB and BM25."""
+    ensemble_retriever = EnsembleRetriever(
+        retrievers=[vector_db.as_retriever(), bm25_retriever],
+        weights=[0.7, 0.3]  # Adjust weights as needed
+    )
+    return ensemble_retriever
+
+# -----------------------------------------------------------------------------
+# Initialize Database
+# -----------------------------------------------------------------------------
+def initialize_database(list_file_obj, progress=gr.Progress()):
+    """Initialize the document database."""
+    list_file_path = [x.name for x in list_file_obj if x is not None]
+    doc_splits = load_doc(list_file_path)
+
+    # Create vector databases and retrievers
+    chromadb = create_chromadb(doc_splits)
+    bm25_retriever = create_bm25_retriever(doc_splits)
+
+    # Create ensemble retriever
+    ensemble_retriever = create_ensemble_retriever(chromadb, bm25_retriever)
 
-def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
-    """Initialize the language model chain"""
+    return ensemble_retriever, "Database created successfully!"
+
+# -----------------------------------------------------------------------------
+# Initialize LLM Chain
+# -----------------------------------------------------------------------------
+def initialize_llmchain(llm_model, temperature, max_tokens, top_k, retriever, progress=gr.Progress()):
+    """Initialize the language model chain."""
     llm = HuggingFaceEndpoint(
         repo_id=llm_model,
         huggingfacehub_api_token=api_token,
@@ -51,14 +135,13 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, pr
         top_k=top_k,
         task="text-generation"
     )
-    
+
     memory = ConversationBufferMemory(
         memory_key="chat_history",
         output_key='answer',
         return_messages=True
     )
 
-    retriever = vector_db.as_retriever()
     qa_chain = ConversationalRetrievalChain.from_llm(
         llm,
         retriever=retriever,
@@ -69,35 +152,50 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, pr
     )
     return qa_chain
 
-def initialize_database(list_file_obj, progress=gr.Progress()):
-    """Initialize the document database"""
-    list_file_path = [x.name for x in list_file_obj if x is not None]
-    doc_splits = load_doc(list_file_path)
-    vector_db = create_db(doc_splits)
-    return vector_db, "Database created successfully!"
-
-def initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
-    """Initialize the Language Model"""
+# -----------------------------------------------------------------------------
+# Initialize LLM
+# -----------------------------------------------------------------------------
+def initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, retriever, progress=gr.Progress()):
+    """Initialize the Language Model."""
     llm_name = list_llm[llm_option]
     print("Selected LLM model:", llm_name)
-    qa_chain = initialize_llmchain(llm_name, llm_temperature, max_tokens, top_k, vector_db, progress)
+    qa_chain = initialize_llmchain(llm_name, llm_temperature, max_tokens, top_k, retriever, progress)
     return qa_chain, "Analysis Assistant initialized and ready!"
 
+# -----------------------------------------------------------------------------
+# Chat History Formatting
+# -----------------------------------------------------------------------------
 def format_chat_history(message, chat_history):
-    """Format chat history for the model"""
+    """Format chat history for the model."""
     formatted_chat_history = []
     for user_message, bot_message in chat_history:
         formatted_chat_history.append(f"User: {user_message}")
         formatted_chat_history.append(f"Assistant: {bot_message}")
     return formatted_chat_history
 
-def conversation(qa_chain, message, history):
-    """Handle conversation and document analysis"""
+# -----------------------------------------------------------------------------
+# Conversation Function
+# -----------------------------------------------------------------------------
+def conversation(qa_chain, message, history, lang):
+    """Handle conversation and document analysis."""
+
+    # Add language instruction to the message
+    if lang == "pt":
+        message += " (Responda em Português)"
+    else:
+        message += " (Respond in English)"
+
     formatted_chat_history = format_chat_history(message, history)
     response = qa_chain.invoke({"question": message, "chat_history": formatted_chat_history})
     response_answer = response["answer"]
+
+    # Remove the language instruction from the chat history
+    if "(Respond" in message:
+        message = message.split(" (Respond")[0]
+
     if response_answer.find("Helpful Answer:") != -1:
         response_answer = response_answer.split("Helpful Answer:")[-1]
+
     response_sources = response["source_documents"]
     response_source1 = response_sources[0].page_content.strip()
     response_source2 = response_sources[1].page_content.strip()
@@ -106,19 +204,20 @@ def conversation(qa_chain, message, history):
     response_source2_page = response_sources[1].metadata["page"] + 1
     response_source3_page = response_sources[2].metadata["page"] + 1
     new_history = history + [(message, response_answer)]
-    return qa_chain, gr.update(value=""), new_history, response_source1, response_source1_page, response_source2, response_source2_page, response_source3, response_source3_page
 
+    return qa_chain, gr.update(value=""), new_history, response_source1, response_source1_page, response_source2, response_source2_page, response_source3, response_source3_page
 
-# ... (código anterior permanece igual)
-
+# -----------------------------------------------------------------------------
+# Gradio Demo
+# -----------------------------------------------------------------------------
 def demo():
-    """Main demo application with enhanced layout"""
+    """Main demo application with enhanced layout."""
     theme = gr.themes.Default(
         primary_hue="indigo",
         secondary_hue="blue",
         neutral_hue="slate",
     )
-    
+
     # Custom CSS for advanced layout
     custom_css = """
         .container {background: #ffffff; padding: 1rem; border-radius: 8px; box-shadow: 0 1px 3px rgba(0,0,0,0.1);}
@@ -129,12 +228,12 @@ def demo():
         .control-panel {margin-bottom: 1rem;}
         .chat-area {background: white; padding: 1rem; border-radius: 8px;}
     """
-    
+
     with gr.Blocks(theme=theme, css=custom_css) as demo:
-        vector_db = gr.State()
+        retriever = gr.State()
         qa_chain = gr.State()
-        language = gr.State(value="en")  # Novo estado para controle de idioma
-        
+        language = gr.State(value="en")  # State for language control
+
         # Header
         gr.HTML(
             """
@@ -144,12 +243,12 @@ def demo():
             </div>
             """
         )
-        
+
         with gr.Row():
             # Left Column - Controls
             with gr.Column(scale=1):
                 gr.Markdown("## Document Processing")
-                
+
                 # File Upload Section
                 with gr.Column(elem_classes="section"):
                     gr.Markdown("### 📄 Upload Documents")
@@ -163,7 +262,7 @@ def demo():
                         value="Ready for documents",
                         label="Processing Status"
                     )
-                
+
                 # Model Selection Section
                 with gr.Column(elem_classes="section"):
                     gr.Markdown("### 🤖 Model Configuration")
@@ -173,15 +272,15 @@ def demo():
                         value=list_llm_simple[0],
                         type="index"
                     )
-                    
-                    # Novo botão de seleção de idioma
+
+                    # Language selection button
                     language_btn = gr.Radio(
                         choices=["English", "Português"],
                         label="Response Language",
                         value="English",
                         type="value"
                     )
-                    
+
                     with gr.Accordion("Advanced Settings", open=False):
                         slider_temperature = gr.Slider(
                             minimum=0.01,
@@ -204,17 +303,17 @@ def demo():
                             step=1,
                             label="Analysis Diversity"
                         )
-                    
+
                     qachain_btn = gr.Button("Initialize Assistant")
                     llm_progress = gr.Textbox(
                         value="Not initialized",
                         label="Assistant Status"
                     )
-            
+
             # Right Column - Chat Interface
             with gr.Column(scale=2):
                 gr.Markdown("## Interactive Analysis")
-                
+
                 # Features Section
                 with gr.Row():
                     with gr.Column():
@@ -235,7 +334,7 @@ def demo():
                             - Specify standards
                             """
                         )
-                
+
                 # Chat Interface
                 with gr.Column(elem_classes="chat-area"):
                     chatbot = gr.Chatbot(
@@ -252,7 +351,7 @@ def demo():
                             [msg, chatbot],
                             value="Clear"
                         )
-                
+
                 # References Section
                 with gr.Accordion("Document References", open=False):
                     with gr.Row():
@@ -271,10 +370,10 @@ def demo():
             """
             ---
             ### About MetroAssist AI
-            
-            A specialized tool for metrology professionals, providing advanced analysis 
+
+            A specialized tool for metrology professionals, providing advanced analysis
             of calibration certificates, measurement data, and technical standards compliance.
-            
+
             **Version 1.0** | © 2024 MetroAssist AI
             """
         )
@@ -285,16 +384,16 @@ def demo():
             inputs=language_btn,
             outputs=language
         )
-        
+
         db_btn.click(
             initialize_database,
             inputs=[document],
-            outputs=[vector_db, db_progress]
+            outputs=[retriever, db_progress]
         )
-        
+
         qachain_btn.click(
             initialize_LLM,
-            inputs=[llm_btn, slider_temperature, slider_maxtokens, slider_topk, vector_db],
+            inputs=[llm_btn, slider_temperature, slider_maxtokens, slider_topk, retriever],
             outputs=[qa_chain, llm_progress]
         ).then(
             lambda: [None, "", 0, "", 0, "", 0],
@@ -309,14 +408,14 @@ def demo():
             outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page],
             queue=False
         )
-        
+
         submit_btn.click(
             conversation,
             inputs=[qa_chain, msg, chatbot, language],
             outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page],
             queue=False
         )
-        
+
         clear_btn.click(
             lambda: [None, "", 0, "", 0, "", 0],
             inputs=None,
@@ -326,36 +425,5 @@ def demo():
 
     demo.queue().launch(debug=True)
 
-# Modifique a função de conversão para incluir o idioma
-def conversation(qa_chain, message, history, lang):
-    """Handle conversation and document analysis"""
-    # Adicione instrução de idioma à mensagem
-    if lang == "pt":
-        message += " (Responda em Português)"
-    else:
-        message += " (Respond in English)"
-    
-    formatted_chat_history = format_chat_history(message, history)
-    response = qa_chain.invoke({"question": message, "chat_history": formatted_chat_history})
-    response_answer = response["answer"]
-    
-    # Remova a instrução de idioma do histórico do chat
-    if "(Respond" in message:
-        message = message.split(" (Respond")[0]
-    
-    if response_answer.find("Helpful Answer:") != -1:
-        response_answer = response_answer.split("Helpful Answer:")[-1]
-    
-    response_sources = response["source_documents"]
-    response_source1 = response_sources[0].page_content.strip()
-    response_source2 = response_sources[1].page_content.strip()
-    response_source3 = response_sources[2].page_content.strip()
-    response_source1_page = response_sources[0].metadata["page"] + 1
-    response_source2_page = response_sources[1].metadata["page"] + 1
-    response_source3_page = response_sources[2].metadata["page"] + 1
-    new_history = history + [(message, response_answer)]
-    
-    return qa_chain, gr.update(value=""), new_history, response_source1, response_source1_page, response_source2, response_source2_page, response_source3, response_source3_page
-
 if __name__ == "__main__":
-    demo()
+    demo()