Update app.py

app.py

@@ -1,128 +1,478 @@
-"""Manual RAG: separate CPU retrieval, CPU fallback, and ZeroGPU generation."""
 
-# MUST come before torch, Gradio, sentence-transformers, or transformers.
-import spaces
+"""
+Document Q&A with Manual RAG
+CPU-Only Gradio Application
+Hugging Face Spaces
+"""
 
 import os
 import traceback
-import torch
+
+# Force CPU execution
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
 import gradio as gr
 
 from rag_core import (
-    load_generator,
     load_retriever,
     answer_question,
     display_evidence,
 )
 
 
-def render_result(result):
-    return result["answer"], display_evidence(result), result["status"]
+# ==========================================
+# CUSTOM UI STYLING
+# ==========================================
 
+CUSTOM_CSS = """
+.gradio-container {
+    max-width: 1080px !important;
+    margin: 0 auto !important;
+    padding: 22px 16px 38px !important;
+}
 
-def respond_retrieval(question):
-    """Retrieval-only processing on CPU; it never asks for a GPU."""
-    try:
-        return render_result(answer_question(question, retrieval_only=True))
-    except Exception as error:
-        traceback.print_exc()
-        return f"Retrieval failed: {error}", "", "CPU retrieval error"
+#rag-hero {
+    padding: 26px 30px;
+    border: 1px solid var(--border-color-primary);
+    border-radius: 20px;
+    background: linear-gradient(
+        125deg,
+        color-mix(
+            in srgb,
+            var(--primary-500) 13%,
+            var(--block-background-fill)
+        ) 0%,
+        var(--block-background-fill) 75%
+    );
+    margin-bottom: 18px;
+}
+
+#rag-hero .eyebrow {
+    font-size: 0.76rem;
+    text-transform: uppercase;
+    font-weight: 750;
+    letter-spacing: 0.11em;
+    color: var(--primary-500);
+    margin-bottom: 9px;
+}
+
+#rag-hero h1 {
+    font-size: clamp(1.65rem, 3vw, 2.35rem);
+    font-weight: 760;
+    line-height: 1.17;
+    letter-spacing: -0.035em;
+    margin: 0 0 10px;
+}
+
+#rag-hero p {
+    color: var(--body-text-color-subdued);
+    font-size: 0.99rem;
+    max-width: 750px;
+    line-height: 1.6;
+    margin: 0;
+}
+
+#rag-hero .meta {
+    font-size: 0.82rem;
+    color: var(--body-text-color-subdued);
+    margin-top: 15px;
+}
+
+#rag-input textarea {
+    font-size: 1.03rem !important;
+}
+
+#rag-actions button {
+    min-height: 48px !important;
+    font-weight: 650;
+}
+
+#rag-status {
+    margin: 4px 0 8px;
+}
+
+#rag-status p {
+    font-size: 0.9rem;
+    color: var(--body-text-color-subdued);
+}
+
+#rag-answer textarea {
+    font-size: 1.04rem !important;
+    line-height: 1.6 !important;
+}
+
+#rag-evidence textarea {
+    font-size: 0.94rem !important;
+    line-height: 1.55 !important;
+}
+
+#rag-footnote p {
+    font-size: 0.82rem;
+    color: var(--body-text-color-subdued);
+}
 
+@media (max-width: 680px) {
+    #rag-hero {
+        padding: 20px 18px;
+        border-radius: 16px;
+    }
+
+    .gradio-container {
+        padding: 12px 10px 22px !important;
+    }
+}
+"""
+
+
+# ==========================================
+# FORMAT RAG RESULTS
+# ==========================================
+
+def format_result(result):
+    """Format answer, evidence, and run status."""
+
+    answer = result.get("answer", "")
+    evidence = display_evidence(result)
+
+    hits = result.get("retrieved") or []
+    count = len(hits)
+
+    top_score = (
+        hits[0]["score"]
+        if hits
+        else None
+    )
+
+    score_note = (
+        f" | Top cosine similarity: {top_score:.3f}"
+        if top_score is not None
+        else ""
+    )
+
+    status = result.get("status", "")
+
+    if status == "retrieval_only":
+
+        answer = (
+            "Evidence retrieved successfully. "
+            "Select Generate Answer to create a response."
+        )
+
+        message = (
+            f"**Evidence retrieved** | "
+            f"{count} source passages{score_note}"
+        )
+
+    elif status == "generated":
+
+        message = (
+            f"**Answer generated on CPU** | "
+            f"{count} source passages{score_note}"
+        )
+
+    elif status == "low_similarity_refusal":
+
+        message = (
+            "**No answer generated** | "
+            "The retrieved evidence is below "
+            "the similarity threshold."
+            f"{score_note}"
+        )
+
+    else:
+
+        message = (
+            f"**Completed on CPU** | "
+            f"{count} source passages{score_note}"
+        )
+
+    return answer, evidence, message
+
+
+# ==========================================
+# RETRIEVE EVIDENCE
+# ==========================================
+
+def retrieve_evidence(question):
+    """Retrieve relevant documents without generation."""
 
-def respond_cpu(question):
-    """CPU answer-generation fallback when ZeroGPU is unavailable."""
     try:
-        return render_result(answer_question(question, device="cpu"))
-    except Exception as error:
+
+        result = answer_question(
+            question,
+            retrieval_only=True
+        )
+
+        return format_result(result)
+
+    except ValueError as error:
+
+        return (
+            "Please check your question.",
+            "",
+            f"**Input error:** {str(error)}"
+        )
+
+    except Exception:
+
         traceback.print_exc()
-        return f"CPU generation failed: {error}", "", "CPU generation error"
 
+        return (
+            "Unable to retrieve evidence.",
+            "",
+            "**Retrieval error.** Check the application logs."
+        )
+
+
+# ==========================================
+# GENERATE ANSWER
+# ==========================================
+
+def generate_answer(question):
+    """Generate an answer using FLAN-T5 on CPU."""
 
-@spaces.GPU(duration=60)
-def respond_gpu(question):
-    """Answer with the model that ZeroGPU has packed for CUDA."""
     try:
-        return render_result(answer_question(question, device="cuda"))
-    except Exception as error:
+
+        result = answer_question(
+            question,
+            retrieval_only=False
+        )
+
+        return format_result(result)
+
+    except ValueError as error:
+
+        return (
+            "Please check your question.",
+            "",
+            f"**Input error:** {str(error)}"
+        )
+
+    except Exception:
+
         traceback.print_exc()
-        return f"GPU generation failed: {error}", "", "GPU generation error"
 
+        return (
+            "Unable to generate an answer. "
+            "You can still retrieve evidence.",
+            "",
+            "**Generation error.** Check the application logs."
+        )
 
-@spaces.GPU(duration=10)
-def check_zerogpu():
-    """Test GPU allocation without loading or invoking any RAG model."""
-    # If ZeroGPU worker_init fails, THIS LINE IS NEVER REACHED.
-    value = torch.ones(1, device="cuda").sum().item()
-    return f"ZeroGPU allocated successfully: {torch.cuda.get_device_name(0)}; test value = {value:.0f}"
 
+# ==========================================
+# BUILD GRADIO INTERFACE
+# ==========================================
 
 def make_demo():
-    with gr.Blocks(title="Document Q&A with Manual RAG") as demo:
-        gr.Markdown(
-            "# Document Q&A with Manual RAG\n"
-            "Ask questions about 30 SQuAD/Wikipedia passages. "
-            "Evidence is retrieved on CPU; you may generate an answer on "
-            "ZeroGPU or use the CPU fallback. Verify answers against sources."
+
+    with gr.Blocks(
+        title="Document Q&A with Manual RAG"
+    ) as demo:
+
+        # Application header
+        gr.HTML(
+            """
+            <section id="rag-hero">
+
+                <div class="eyebrow">
+                    Manual RAG | CPU Powered
+                </div>
+
+                <h1>
+                    Document Q&amp;A with Manual RAG
+                </h1>
+
+                <p>
+                    Ask a question, retrieve supporting
+                    passages, and generate an answer
+                    using only the available documents.
+                </p>
+
+                <div class="meta">
+                    30 SQuAD/Wikipedia passages
+                    &nbsp; | &nbsp;
+                    MiniLM Retrieval
+                    &nbsp; | &nbsp;
+                    FLAN-T5 Small Generation
+                </div>
+
+            </section>
+            """
         )
+
+        # Question input
         question = gr.Textbox(
-            label="Question",
+            label="Ask a Question",
             placeholder="Which company produces the iPod?",
             lines=2,
+            max_lines=3,
+            max_length=300,
+            elem_id="rag-input"
         )
-        with gr.Row():
-            retrieve_button = gr.Button("Retrieve evidence (CPU)")
-            gpu_button = gr.Button("Answer with ZeroGPU", variant="primary")
-            cpu_button = gr.Button("Answer with CPU fallback")
-        answer = gr.Textbox(label="Answer - verify against evidence", lines=5)
-        evidence = gr.Textbox(label="Source passages and similarity scores", lines=14)
-        status = gr.Textbox(label="Run status")
-        with gr.Accordion("ZeroGPU allocation diagnostic", open=False):
-            gr.Markdown(
-                "This runs one tiny CUDA operation. If it fails with "
-                "`No CUDA GPUs are available` in `worker_init`, the error "
-                "occurs before your RAG function runs."
+
+        # Action buttons
+        with gr.Row(elem_id="rag-actions"):
+
+            answer_button = gr.Button(
+                "Generate Answer",
+                variant="primary",
+                size="lg"
+            )
+
+            retrieve_button = gr.Button(
+                "Retrieve Evidence",
+                variant="secondary",
+                size="lg"
             )
-            test_button = gr.Button("Test ZeroGPU allocation")
-            gpu_test_result = gr.Textbox(label="GPU test result")
-
-        outputs = [answer, evidence, status]
-        retrieve_button.click(fn=respond_retrieval, inputs=question, outputs=outputs)
-        cpu_button.click(fn=respond_cpu, inputs=question, outputs=outputs)
-        gpu_button.click(fn=respond_gpu, inputs=question, outputs=outputs)
-        test_button.click(fn=check_zerogpu, inputs=[], outputs=gpu_test_result)
+
+        # Application status
+        status = gr.Markdown(
+            "**Ready** | Running on CPU.",
+            elem_id="rag-status"
+        )
+
+        # Answer display
+        gr.Markdown("### Generated Answer")
+
+        answer = gr.Textbox(
+            label="Answer",
+            placeholder="Your answer will appear here.",
+            lines=4,
+            max_lines=8,
+            interactive=False,
+            buttons=["copy"],
+            elem_id="rag-answer"
+        )
+
+        # Retrieved documents
+        with gr.Accordion(
+            "Source Passages and Similarity Scores",
+            open=True
+        ):
+
+            evidence = gr.Textbox(
+                show_label=False,
+                placeholder=(
+                    "Retrieved evidence will appear here."
+                ),
+                lines=11,
+                max_lines=14,
+                interactive=False,
+                buttons=["copy"],
+                elem_id="rag-evidence"
+            )
+
+        # Example questions
         gr.Examples(
             examples=[
-                ["Which company produces the iPod?"],
-                ["At what age did Frédéric move to Paris?"],
-                ["What is Beyoncé's private mobile phone number?"],
+                [
+                    "Which company produces the iPod?"
+                ],
+                [
+                    "At what age did Frédéric move to Paris?"
+                ],
+                [
+                    "What is the live price of Bitcoin "
+                    "in Nigerian naira right now?"
+                ]
             ],
             inputs=[question],
+            label="Try an Example Question",
+            cache_examples=False
         )
+
+        # Footer
+        gr.Markdown(
+            """
+            **Important:** Answers are generated from
+            the 30 available SQuAD/Wikipedia passages,
+            not the live internet.
+
+            Similarity scores are not confidence scores.
+            Always verify answers against the retrieved
+            source passages.
+            """,
+            elem_id="rag-footnote"
+        )
+
+        # Connect buttons
+        outputs = [
+            answer,
+            evidence,
+            status
+        ]
+
+        answer_button.click(
+            fn=generate_answer,
+            inputs=question,
+            outputs=outputs
+        )
+
+        retrieve_button.click(
+            fn=retrieve_evidence,
+            inputs=question,
+            outputs=outputs
+        )
+
     return demo
 
 
-# Run initialization on the MAIN THREAD at module import; no model loads in
-# Gradio callbacks. In ZeroGPU's parent process, CUDA must remain uninitialized.
-print("Starting Manual RAG (ZeroGPU diagnostics)", flush=True)
-print(
-    f"ZeroGPU env: {os.getenv('SPACES_ZERO_GPU')} | "
-    f"torch={torch.__version__} | "
-    f"parent CUDA initialized={torch.cuda.is_initialized()}",
-    flush=True,
-)
-print("Preloading FLAN-T5 onto ZeroGPU-managed CUDA...", flush=True)
-print("Generation model:", load_generator(), flush=True)
-print("Loading CPU retriever:", load_retriever(), flush=True)
-print("Parent CUDA initialized after load:", torch.cuda.is_initialized(), flush=True)
+# ==========================================
+# START APPLICATION
+# ==========================================
 
-demo = make_demo()
-demo.queue(default_concurrency_limit=1, max_size=20)
+def main():
 
-if __name__ == "__main__":
+    print("=" * 50)
+    print("Starting Document Q&A with Manual RAG")
+    print("Execution Mode: CPU Only")
+    print("=" * 50)
+
+    try:
+
+        print("Loading document retriever...")
+
+        info = load_retriever()
+
+        print(
+            f"Retriever initialized successfully: {info}"
+        )
+
+    except Exception:
+
+        print("Failed to initialize retriever.")
+        traceback.print_exc()
+        raise
+
+    # Build UI
+    demo = make_demo()
+
+    # Configure queue
+    demo.queue(
+        default_concurrency_limit=1,
+        max_size=20
+    )
+
+    print("Launching Gradio application...")
+
+    # Launch the app
     demo.launch(
         server_name="0.0.0.0",
-        server_port=int(os.getenv("PORT", "7860")),
-        show_error=True,
+        server_port=int(
+            os.environ.get("PORT", "7860")
+        ),
         share=False,
+        show_error=True,
         ssr_mode=False,
+        theme=gr.themes.Soft(
+            primary_hue="indigo",
+            neutral_hue="slate"
+        ),
+        css=CUSTOM_CSS
     )
+
+
+if __name__ == "__main__":
+    main()