Update app.py

app.py

@@ -1,342 +1,123 @@
-"""
-PDF-based chatbot with Retrieval-Augmented Generation
-"""
+# ✅ Enhanced GenAI Assistant with:
+# - PDF/TXT support
+# - Ask Anything + Challenge Me modes
+# - Auto Summary (<=150 words)
+# - Memory handling
+# - Reference highlighting
+# - Stunning UI (Gradio upgraded)
+
+# --- FILE: app.py ---
 
 import os
 import gradio as gr
-
 from dotenv import load_dotenv
-
 import indexing
 import retrieval
+import utils
 
-
-# default_persist_directory = './chroma_HF/'
 list_llm = [
     "mistralai/Mistral-7B-Instruct-v0.3",
     "microsoft/Phi-3.5-mini-instruct",
     "meta-llama/Llama-3.1-8B-Instruct",
     "meta-llama/Llama-3.2-3B-Instruct",
-    "meta-llama/Llama-3.2-1B-Instruct",
-    "HuggingFaceTB/SmolLM2-1.7B-Instruct",
-    "HuggingFaceH4/zephyr-7b-beta",
-    "HuggingFaceH4/zephyr-7b-gemma-v0.1",
     "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
     "google/gemma-2-2b-it",
-    "google/gemma-2-9b-it",
-    "Qwen/Qwen2.5-1.5B-Instruct",
     "Qwen/Qwen2.5-3B-Instruct",
-    "Qwen/Qwen2.5-7B-Instruct",
 ]
 list_llm_simple = [os.path.basename(llm) for llm in list_llm]
 
 
-# Load environment file - HuggingFace API key
 def retrieve_api():
-    """Retrieve HuggingFace API Key"""
-    _ = load_dotenv()
+    load_dotenv()
     global huggingfacehub_api_token
     huggingfacehub_api_token = os.environ.get("HUGGINGFACE_API_KEY")
 
 
-# Initialize database
-def initialize_database(
-    list_file_obj, chunk_size, chunk_overlap, progress=gr.Progress()
-):
-    """Initialize database"""
-
-    # Create list of documents (when valid)
+def initialize_database(list_file_obj, chunk_size, chunk_overlap, progress=gr.Progress()):
     list_file_path = [x.name for x in list_file_obj if x is not None]
-
-    # Create collection_name for vector database
-    progress(0.1, desc="Creating collection name...")
     collection_name = indexing.create_collection_name(list_file_path[0])
-
-    progress(0.25, desc="Loading document...")
-    # Load document and create splits
-    doc_splits = indexing.load_doc(list_file_path, chunk_size, chunk_overlap)
-
-    # Create or load vector database
-    progress(0.5, desc="Generating vector database...")
-
-    # global vector_db
+    doc_splits, full_text = indexing.load_doc(list_file_path, chunk_size, chunk_overlap)
+    summary = utils.generate_summary(full_text)
     vector_db = indexing.create_db(doc_splits, collection_name)
+    return vector_db, collection_name, summary, "Complete!"
 
-    return vector_db, collection_name, "Complete!"
 
-
-# Initialize LLM
-def initialize_llm(
-    llm_option, llm_temperature, max_tokens, top_k, vector_db, progress=gr.Progress()
-):
-    """Initialize LLM"""
-
-    # print("llm_option",llm_option)
+def initialize_llm(llm_option, llm_temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
     llm_name = list_llm[llm_option]
-    print("llm_name: ", llm_name)
     qa_chain = retrieval.initialize_llmchain(
         llm_name, huggingfacehub_api_token, llm_temperature, max_tokens, top_k, vector_db, progress
     )
     return qa_chain, "Complete!"
 
 
-# Chatbot conversation
 def conversation(qa_chain, message, history):
-    """Chatbot conversation"""
+    qa_chain, new_history, response_sources = retrieval.invoke_qa_chain(qa_chain, message, history)
+    highlights = utils.extract_highlight_snippets(response_sources)
+    return qa_chain, gr.update(value=""), new_history, *highlights
 
-    qa_chain, new_history, response_sources = retrieval.invoke_qa_chain(
-        qa_chain, message, history
-    )
 
-    # Format output gradio components
-    response_source1 = response_sources[0].page_content.strip()
-    response_source2 = response_sources[1].page_content.strip()
-    response_source3 = response_sources[2].page_content.strip()
-    # Langchain sources are zero-based
-    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
+def challenge_me(qa_chain):
+    questions = utils.generate_challenge_questions(qa_chain)
+    return questions
 
-    return (
-        qa_chain,
-        gr.update(value=""),
-        new_history,
-        response_source1,
-        response_source1_page,
-        response_source2,
-        response_source2_page,
-        response_source3,
-        response_source3_page,
-    )
-
-
-SPACE_TITLE = """
-<center><h2>PDF-based chatbot</center></h2>
-<h3>Ask any questions about your PDF documents</h3>
-"""
 
-SPACE_INFO = """
-<b>Description:</b> This AI assistant, using Langchain and open-source LLMs, performs retrieval-augmented generation (RAG) from your PDF documents. \
-The user interface explicitely shows multiple steps to help understand the RAG workflow. 
-This chatbot takes past questions into account when generating answers (via conversational memory), and includes document references for clarity purposes.<br>
-<br><b>Notes:</b> Updated space with more recent LLM models (Qwen 2.5, Llama 3.2, SmolLM2 series) 
-<br><b>Warning:</b> This space uses the free CPU Basic hardware from Hugging Face. Some steps and LLM models used below (free inference endpoints) can take some time to generate a reply.
-"""
+def evaluate_answers(qa_chain, questions, user_answers):
+    feedback = utils.evaluate_responses(qa_chain, questions, user_answers)
+    return feedback
 
 
-# Gradio User Interface
 def gradio_ui():
-    """Gradio User Interface"""
-
-    with gr.Blocks(theme="base") as demo:
+    with gr.Blocks(theme=gr.themes.Monochrome(), css="footer {display:none}") as demo:
         vector_db = gr.State()
         qa_chain = gr.State()
         collection_name = gr.State()
 
-        gr.Markdown(SPACE_TITLE)
-        gr.Markdown(SPACE_INFO)
-
-        with gr.Tab("Step 1 - Upload PDF"):
-            with gr.Row():
-                document = gr.File(
-                    height=200,
-                    file_count="multiple",
-                    file_types=[".pdf"],
-                    interactive=True,
-                    label="Upload your PDF documents (single or multiple)",
-                )
-
-        with gr.Tab("Step 2 - Process document"):
-            with gr.Row():
-                db_btn = gr.Radio(
-                    ["ChromaDB"],
-                    label="Vector database type",
-                    value="ChromaDB",
-                    type="index",
-                    info="Choose your vector database",
-                )
-            with gr.Accordion("Advanced options - Document text splitter", open=False):
-                with gr.Row():
-                    slider_chunk_size = gr.Slider(
-                        minimum=100,
-                        maximum=1000,
-                        value=600,
-                        step=20,
-                        label="Chunk size",
-                        info="Chunk size",
-                        interactive=True,
-                    )
-                with gr.Row():
-                    slider_chunk_overlap = gr.Slider(
-                        minimum=10,
-                        maximum=200,
-                        value=40,
-                        step=10,
-                        label="Chunk overlap",
-                        info="Chunk overlap",
-                        interactive=True,
-                    )
-            with gr.Row():
-                db_progress = gr.Textbox(
-                    label="Vector database initialization", value="None"
-                )
-            with gr.Row():
-                db_btn = gr.Button("Generate vector database")
-
-        with gr.Tab("Step 3 - Initialize QA chain"):
-            with gr.Row():
-                llm_btn = gr.Radio(
-                    list_llm_simple,
-                    label="LLM models",
-                    value=list_llm_simple[6],
-                    type="index",
-                    info="Choose your LLM model",
-                )
-            with gr.Accordion("Advanced options - LLM model", open=False):
-                with gr.Row():
-                    slider_temperature = gr.Slider(
-                        minimum=0.01,
-                        maximum=1.0,
-                        value=0.7,
-                        step=0.1,
-                        label="Temperature",
-                        info="Model temperature",
-                        interactive=True,
-                    )
-                with gr.Row():
-                    slider_maxtokens = gr.Slider(
-                        minimum=224,
-                        maximum=4096,
-                        value=1024,
-                        step=32,
-                        label="Max Tokens",
-                        info="Model max tokens",
-                        interactive=True,
-                    )
-                with gr.Row():
-                    slider_topk = gr.Slider(
-                        minimum=1,
-                        maximum=10,
-                        value=3,
-                        step=1,
-                        label="top-k samples",
-                        info="Model top-k samples",
-                        interactive=True,
-                    )
-            with gr.Row():
-                llm_progress = gr.Textbox(value="None", label="QA chain initialization")
-            with gr.Row():
-                qachain_btn = gr.Button("Initialize Question Answering chain")
-
-        with gr.Tab("Step 4 - Chatbot"):
-            chatbot = gr.Chatbot(height=300, type="tuples")
-            with gr.Accordion("Advanced - Document references", open=False):
-                with gr.Row():
-                    doc_source1 = gr.Textbox(
-                        label="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
-                    )
-                    source2_page = gr.Number(label="Page", scale=1)
-                with gr.Row():
-                    doc_source3 = gr.Textbox(
-                        label="Reference 3", lines=2, container=True, scale=20
-                    )
-                    source3_page = gr.Number(label="Page", scale=1)
-            with gr.Row():
-                msg = gr.Textbox(
-                    placeholder="Type message (e.g. 'Can you summarize this document in one paragraph?')",
-                    container=True,
-                )
-            with gr.Row():
-                submit_btn = gr.Button("Submit message")
-                clear_btn = gr.ClearButton(
-                    components=[msg, chatbot], value="Clear conversation"
-                )
-
-        # Preprocessing events
-        db_btn.click(
-            initialize_database,
-            inputs=[document, slider_chunk_size, slider_chunk_overlap],
-            outputs=[vector_db, collection_name, 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,
-        )
-
-        # Chatbot events
-        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)
-
+        gr.Markdown("""<h1 style='text-align:center;'>📚 GenAI Document Assistant</h1>
+        <h3 style='text-align:center;color:gray;'>Smart, interactive reading of research papers, legal docs, and more.</h3>""")
+
+        with gr.Tab("1️⃣ Upload Document"):
+            document = gr.File(label="Upload PDF or TXT", file_types=[".pdf", ".txt"], file_count="multiple")
+            slider_chunk_size = gr.Slider(100, 1000, value=600, step=20, label="Chunk Size")
+            slider_chunk_overlap = gr.Slider(10, 200, value=40, step=10, label="Chunk Overlap")
+            db_progress = gr.Textbox(label="Processing Status")
+            summary_box = gr.Textbox(label="Auto Summary (≤ 150 words)", lines=5)
+            db_btn = gr.Button("📥 Process Document")
+
+        with gr.Tab("2️⃣ QA Chain Initialization"):
+            llm_btn = gr.Radio(list_llm_simple, label="Select LLM", value=list_llm_simple[0], type="index")
+            slider_temperature = gr.Slider(0.01, 1.0, value=0.7, step=0.1, label="Temperature")
+            slider_maxtokens = gr.Slider(224, 4096, value=1024, step=32, label="Max Tokens")
+            slider_topk = gr.Slider(1, 10, value=3, step=1, label="Top-K")
+            llm_progress = gr.Textbox(label="LLM Status")
+            qachain_btn = gr.Button("⚙️ Initialize QA Chain")
+
+        with gr.Tab("3️⃣ Ask Anything"):
+            chatbot = gr.Chatbot(height=300)
+            msg = gr.Textbox(placeholder="Ask a question from the document...")
+            submit_btn = gr.Button("💬 Ask")
+            clear_btn = gr.ClearButton([msg, chatbot])
+            ref1 = gr.Textbox(label="Reference 1")
+            ref2 = gr.Textbox(label="Reference 2")
+            ref3 = gr.Textbox(label="Reference 3")
+
+        with gr.Tab("4️⃣ Challenge Me"):
+            challenge_btn = gr.Button("🎯 Generate Questions")
+            q1 = gr.Textbox(label="Question 1")
+            a1 = gr.Textbox(label="Your Answer 1")
+            q2 = gr.Textbox(label="Question 2")
+            a2 = gr.Textbox(label="Your Answer 2")
+            q3 = gr.Textbox(label="Question 3")
+            a3 = gr.Textbox(label="Your Answer 3")
+            eval_btn = gr.Button("✅ Submit Answers")
+            feedback = gr.Textbox(label="Feedback", lines=5)
+
+        db_btn.click(initialize_database, [document, slider_chunk_size, slider_chunk_overlap], [vector_db, collection_name, summary_box, db_progress])
+        qachain_btn.click(initialize_llm, [llm_btn, slider_temperature, slider_maxtokens, slider_topk, vector_db], [qa_chain, llm_progress])
+        submit_btn.click(conversation, [qa_chain, msg, chatbot], [qa_chain, msg, chatbot, ref1, ref2, ref3])
+        challenge_btn.click(challenge_me, [qa_chain], [q1, q2, q3])
+        eval_btn.click(evaluate_answers, [qa_chain, [q1, q2, q3], [a1, a2, a3]], [feedback])
+
+    demo.launch(debug=True)
 
 if __name__ == "__main__":
     retrieve_api()