Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 import os
-api_token = os.getenv("HF_TOKEN")
-
+import torch
 from langchain_community.vectorstores import FAISS
 from langchain_community.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
@@ -12,9 +11,10 @@ from langchain_community.llms import HuggingFacePipeline
 from langchain.chains import ConversationChain
 from langchain.memory import ConversationBufferMemory
 from langchain_community.llms import HuggingFaceEndpoint
-import torch
 
-# Added Deepseek model to the list
+api_token = os.getenv("HF_TOKEN")
+
+# Available LLM models
 list_llm = [
     "meta-llama/Meta-Llama-3-8B-Instruct", 
     "mistralai/Mistral-7B-Instruct-v0.2",
@@ -22,25 +22,33 @@ list_llm = [
 ]  
 list_llm_simple = [os.path.basename(llm) for llm in list_llm]
 
-# Rest of the functions remain the same until demo()
 def load_doc(list_file_path):
+    """
+    Load and split PDF documents into chunks
+    """
     loaders = [PyPDFLoader(x) for x in list_file_path]
     pages = []
     for loader in loaders:
         pages.extend(loader.load())
     text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size = 1024, 
-        chunk_overlap = 64 
-    )  
+        chunk_size=1024,
+        chunk_overlap=64
+    )
     doc_splits = text_splitter.split_documents(pages)
     return doc_splits
 
 def create_db(splits):
+    """
+    Create vector database from document splits
+    """
     embeddings = HuggingFaceEmbeddings()
     vectordb = FAISS.from_documents(splits, embeddings)
     return vectordb
 
 def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
+    """
+    Initialize the language model chain
+    """
     llm = HuggingFaceEndpoint(
         repo_id=llm_model,
         huggingfacehub_api_token=api_token,
@@ -60,7 +68,7 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, pr
     qa_chain = ConversationalRetrievalChain.from_llm(
         llm,
         retriever=retriever,
-        chain_type="stuff", 
+        chain_type="stuff",
         memory=memory,
         return_source_documents=True,
         verbose=False,
@@ -68,18 +76,27 @@ 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!"
+    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
+    """
     llm_name = list_llm[llm_option]
-    print("llm_name: ",llm_name)
+    print("Selected LLM model:", llm_name)
     qa_chain = initialize_llmchain(llm_name, llm_temperature, max_tokens, top_k, vector_db, progress)
-    return qa_chain, "QA chain initialized. Chatbot is ready!"
+    return qa_chain, "Analysis Assistant initialized and ready!"
 
 def format_chat_history(message, chat_history):
+    """
+    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}")
@@ -87,6 +104,9 @@ def format_chat_history(message, chat_history):
     return formatted_chat_history
 
 def conversation(qa_chain, message, history):
+    """
+    Handle conversation and document analysis
+    """
     formatted_chat_history = format_chat_history(message, history)
     response = qa_chain.invoke({"question": message, "chat_history": formatted_chat_history})
     response_answer = response["answer"]
@@ -102,95 +122,272 @@ def conversation(qa_chain, message, history):
     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
 
-def upload_file(file_obj):
-    list_file_path = []
-    for idx, file in enumerate(file_obj):
-        file_path = file_obj.name
-        list_file_path.append(file_path)
-    return list_file_path
-
 def demo():
-    # Modified theme to use dark blue colors
+    """
+    Main demo application
+    """
+    # Enhanced theme with professional colors
     theme = gr.themes.Default(
         primary_hue="indigo",
         secondary_hue="blue",
-        neutral_hue="slate"
+        neutral_hue="slate",
+        font=[gr.themes.GoogleFont("Roboto"), "system-ui", "sans-serif"]
     )
     
-    with gr.Blocks(theme=theme) as demo:
+    css = """
+        .container { max-width: 1200px; margin: auto; }
+        .metadata { font-size: 0.9em; color: #666; }
+        .highlight { background-color: #f0f7ff; padding: 1em; border-radius: 8px; }
+        .warning { color: #e53e3e; }
+        .success { color: #38a169; }
+    """
+    
+    with gr.Blocks(theme=theme, css=css) as demo:
         vector_db = gr.State()
         qa_chain = gr.State()
-        gr.HTML("<center><h1>RAG PDF chatbot</h1><center>")
-        gr.Markdown("""<b>Query your PDF documents!</b> This AI agent is designed to perform retrieval augmented generation (RAG) on PDF documents. The app is hosted on Hugging Face Hub for the sole purpose of demonstration. \
-        <b>Please do not upload confidential documents.</b>
-        """)
+        
+        # Enhanced header
+        gr.HTML(
+            """
+            <div style='text-align: center; padding: 20px;'>
+                <h1 style='color: #1a365d; margin-bottom: 10px;'>MetroAssist AI - Expert in Metrology Report Analysis</h1>
+                <p style='color: #4a5568; font-size: 1.2em;'>Your intelligent assistant for advanced analysis of metrological documents</p>
+            </div>
+            """
+        )
+        
+        # Marketing and feature description
+        gr.Markdown(
+            """
+            ### 🔍 Specialized Metrology Analysis
+            
+            MetroAssist AI is a specialized assistant designed to revolutionize metrology report analysis. 
+            Powered by cutting-edge AI technology, it offers:
+
+            * **Precise Analysis**: Detailed interpretation of measurements, calibrations, and compliance
+            * **Intelligent Contextualization**: Deep understanding of metrological standards and norms
+            * **Advanced Technical Support**: Assistance in complex instrument and measurement analyses
+            * **Rapid Processing**: Efficient analysis of multiple technical documents
+            
+            ⚠️ **Security Note**: Your documents are processed with complete security. We do not permanently store confidential data.
+            """
+        )
+
         with gr.Row():
-            with gr.Column(scale = 86):
-                gr.Markdown("<b>Step 1 - Upload PDF documents and Initialize RAG pipeline</b>")
+            with gr.Column(scale=86):
+                gr.Markdown(
+                    """
+                    ### 📥 Step 1: Document Loading and Preparation
+                    Upload your metrology reports for expert analysis.
+                    """
+                )
                 with gr.Row():
-                    document = gr.Files(height=300, file_count="multiple", file_types=["pdf"], interactive=True, label="Upload PDF documents")
+                    document = gr.Files(
+                        height=300,
+                        file_count="multiple",
+                        file_types=["pdf"],
+                        interactive=True,
+                        label="Upload Metrology Reports (PDF)",
+                        info="Accepts multiple PDF files"
+                    )
                 with gr.Row():
-                    db_btn = gr.Button("Create vector database")
+                    db_btn = gr.Button(
+                        "Process Documents",
+                        variant="primary",
+                        size="lg"
+                    )
                 with gr.Row():
-                    db_progress = gr.Textbox(value="Not initialized", show_label=False)
-                gr.Markdown("<style>body { font-size: 16px; }</style><b>Select Large Language Model (LLM) and input parameters</b>")
+                    db_progress = gr.Textbox(
+                        value="Waiting for documents...",
+                        show_label=False,
+                        container=False
+                    )
+                
+                gr.Markdown(
+                    """
+                    ### 🤖 Analysis Engine Configuration
+                    Select and configure the AI model to best meet your needs.
+                    """
+                )
                 with gr.Row():
-                    llm_btn = gr.Radio(list_llm_simple, label="Available LLMs", value = list_llm_simple[0], type="index")
+                    llm_btn = gr.Radio(
+                        list_llm_simple,
+                        label="Available AI Models",
+                        value=list_llm_simple[0],
+                        type="index",
+                        info="Choose the most suitable model for your analysis"
+                    )
+                
                 with gr.Row():
-                    with gr.Accordion("LLM input parameters", open=False):
+                    with gr.Accordion("Advanced Analysis Parameters", open=False):
                         with gr.Row():
-                            slider_temperature = gr.Slider(minimum = 0.01, maximum = 1.0, value=0.5, step=0.1, label="Temperature", info="Controls randomness in token generation", interactive=True)
+                            slider_temperature = gr.Slider(
+                                minimum=0.01,
+                                maximum=1.0,
+                                value=0.5,
+                                step=0.1,
+                                label="Analysis Precision",
+                                info="Controls the balance between precision and creativity in analysis",
+                                interactive=True
+                            )
                         with gr.Row():
-                            slider_maxtokens = gr.Slider(minimum = 128, maximum = 9192, value=4096, step=128, label="Max New Tokens", info="Maximum number of tokens to be generated",interactive=True)
+                            slider_maxtokens = gr.Slider(
+                                minimum=128,
+                                maximum=9192,
+                                value=4096,
+                                step=128,
+                                label="Response Length",
+                                info="Defines the level of detail in analyses",
+                                interactive=True
+                            )
                         with gr.Row():
-                            slider_topk = gr.Slider(minimum = 1, maximum = 10, value=3, step=1, label="top-k", info="Number of tokens to select the next token from", interactive=True)
+                            slider_topk = gr.Slider(
+                                minimum=1,
+                                maximum=10,
+                                value=3,
+                                step=1,
+                                label="Analysis Diversity",
+                                info="Controls the variety of perspectives in analysis",
+                                interactive=True
+                            )
+                
                 with gr.Row():
-                    qachain_btn = gr.Button("Initialize Question Answering Chatbot")
+                    qachain_btn = gr.Button(
+                        "Initialize Analysis Assistant",
+                        variant="primary",
+                        size="lg"
+                    )
                 with gr.Row():
-                    llm_progress = gr.Textbox(value="Not initialized", show_label=False)
+                    llm_progress = gr.Textbox(
+                        value="Waiting for initialization...",
+                        show_label=False
+                    )
 
-            with gr.Column(scale = 200):
-                gr.Markdown("<b>Step 2 - Chat with your Document</b>")
-                chatbot = gr.Chatbot(height=505)
-                with gr.Accordion("Relevent context from the source document", open=False):
+            with gr.Column(scale=200):
+                gr.Markdown(
+                    """
+                    ### 💬 Step 2: Expert Consultation and Analysis
+                    Ask questions about your metrology reports. The assistant will provide detailed technical analyses.
+                    
+                    **Suggested questions:**
+                    - Analyze the calibration results of this instrument
+                    - Verify compliance with technical standards
+                    - Identify critical points in measurements
+                    - Compare results with specified limits
+                    - Evaluate measurement uncertainty
+                    - Assess calibration intervals
+                    """
+                )
+                chatbot = gr.Chatbot(
+                    height=505,
+                    show_label=True,
+                    container=True,
+                    label="Metrology Analysis"
+                )
+                
+                with gr.Accordion("Source Document References", open=False):
                     with gr.Row():
-                        doc_source1 = gr.Textbox(label="Reference 1", lines=2, container=True, scale=20)
+                        doc_source1 = gr.Textbox(
+                            label="Technical Reference 1",
+                            lines=2,
+                            container=True,
+                            scale=20
+                        )
                         source1_page = gr.Number(label="Page", scale=1)
                     with gr.Row():
-                        doc_source2 = gr.Textbox(label="Reference 2", lines=2, container=True, scale=20)
+                        doc_source2 = gr.Textbox(
+                            label="Technical Reference 2",
+                            lines=2,
+                            container=True,
+                            scale=20
+                        )
                         source2_page = gr.Number(label="Page", scale=1)
                     with gr.Row():
-                        doc_source3 = gr.Textbox(label="Reference 3", lines=2, container=True, scale=20)
+                        doc_source3 = gr.Textbox(
+                            label="Technical Reference 3",
+                            lines=2,
+                            container=True,
+                            scale=20
+                        )
                         source3_page = gr.Number(label="Page", scale=1)
+                
                 with gr.Row():
-                    msg = gr.Textbox(placeholder="Ask a question", container=True)
+                    msg = gr.Textbox(
+                        placeholder="Enter your question about the metrology report...",
+                        container=True,
+                        label="Your Query"
+                    )
                 with gr.Row():
-                    submit_btn = gr.Button("Submit")
-                    clear_btn = gr.ClearButton([msg, chatbot], value="Clear")
+                    submit_btn = gr.Button(
+                        "Submit Query",
+                        variant="primary"
+                    )
+                    clear_btn = gr.ClearButton(
+                        [msg, chatbot],
+                        value="Clear Conversation",
+                        variant="secondary"
+                    )
+
+        # Footer
+        gr.Markdown(
+            """
+            ---
+            ### ℹ️ About MetroAssist AI
             
+            Developed for metrology professionals, engineers, and technicians who need precise
+            and reliable analysis of technical documents. Our tool uses advanced AI technology
+            to provide valuable insights and support decision-making in metrology.
+            
+            **Specialized Features:**
+            - Detailed analysis of calibration certificates
+            - Interpretation of complex metrological data
+            - Verification of compliance with technical standards
+            - Decision support in metrological processes
+            - Uncertainty analysis and measurement traceability
+            - Quality control and measurement system analysis
+            
+            *Version 1.0 - Updated 2024*
+            """
+        )
+
         # Event handlers
-        db_btn.click(initialize_database, 
-            inputs=[document], 
-            outputs=[vector_db, db_progress])
-        qachain_btn.click(initialize_LLM, 
-            inputs=[llm_btn, slider_temperature, slider_maxtokens, slider_topk, vector_db], 
-            outputs=[qa_chain, llm_progress]).then(lambda:[None,"",0,"",0,"",0], 
-            inputs=None, 
-            outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], 
-            queue=False)
-
-        msg.submit(conversation, 
-            inputs=[qa_chain, msg, chatbot], 
-            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], 
-            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, 
-            outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], 
-            queue=False)
+        db_btn.click(
+            initialize_database,
+            inputs=[document],
+            outputs=[vector_db, db_progress]
+        )
+        
+        qachain_btn.click(
+            initialize_LLM,
+            inputs=[llm_btn, slider_temperature, slider_maxtokens, slider_topk, vector_db],
+            outputs=[qa_chain, llm_progress]
+        ).then(
+            lambda: [None, "", 0, "", 0, "", 0],
+            inputs=None,
+            outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page],
+            queue=False
+        )
+
+        msg.submit(
+            conversation,
+            inputs=[qa_chain, msg, chatbot],
+            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],
+            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,
+            outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page],
+            queue=False
+        )
 
     demo.queue().launch(debug=True)