Upload app.py

app.py

@@ -0,0 +1,403 @@
+import html  # Added for escaping HTML
+import json
+import logging  # Added for status check logging
+import os  # Added for environment variables
+
+import gradio as gr
+import numpy as np
+from sentence_transformers import SentenceTransformer
+
+# Added HfApi for endpoint check
+from huggingface_hub import InferenceClient
+from dotenv import load_dotenv  # Added for .env loading
+
+import utils.interface_utils as interface_utils
+import utils.llm_utils as llm_utils
+
+# Load environment variables from .env file
+load_dotenv()
+
+# REMOVED Endpoint name constant (will be in llm_utils)
+# LLM_ENDPOINT_NAME = "phi-4-max"
+
+# REMOVED Endpoint Status Check Function
+
+# --- Load Data and Models ---
+# These should be loaded once when the app starts.
+print("Loading data and models...")
+
+# Load data from files
+processed_docs = json.load(
+    open("docs_passages_storage/processed_docs.json", encoding="utf-8")
+)
+passages = json.load(open("docs_passages_storage/passages.json", encoding="utf-8"))
+doc_embeds = np.load("docs_passages_storage/passage_embeddings.npy")
+
+# Load the embedding model - Force CPU
+print("Loading embedding model on CPU...")
+model_name = "Snowflake/snowflake-arctic-embed-l-v2.0"
+embed_model = SentenceTransformer(model_name, device="cpu")
+print("Embedding model loaded.")
+
+# Initialize HF Client globally using env vars
+hf_api_token = os.getenv("HF_TOKEN")
+if not hf_api_token:
+    print(
+        "Warning: HF_TOKEN environment variable not set. Inference client might fail."
+    )
+hf_client = InferenceClient(token=hf_api_token)
+
+print("Data and models loaded.")
+
+
+# --- Main Gradio Function ---
+def get_results(query: str, progress=gr.Progress(track_tqdm=True)) -> list:
+    """Processes query, retrieves passages, processes each, formats output, tracks progress."""
+    # Define placeholders for 10 outputs (9 results + 1 summary)
+    if not query:
+        # Return default values for all 10 output slots
+        return ["-"] * 9 + ["### Summary\n-"]
+
+    # --- Check Endpoint Status FIRST ---
+    # Use token loaded previously for the client
+    endpoint_status = llm_utils.check_endpoint_status(
+        token=hf_api_token
+    )  # Call function from llm_utils
+    if endpoint_status["status"] == "error":
+        error_message = endpoint_status["ui_message"]
+        logging.error(f"Endpoint status error: {error_message}")
+        progress(1, desc="Endpoint Error")
+        # Display error in first result slot and summary slot
+        return [
+            "### Endpoint Error",
+            f"<p style='color: red;'>{html.escape(error_message)}</p>",
+            "_",
+            "-",
+            "-",
+            "-",
+            "-",
+            "-",
+            "-",
+            "### Summary\n_Endpoint not ready._",
+        ]
+
+    print(f"Processing query: {query}")
+    try:
+        progress(0, desc="Retrieving relevant documents...")
+        # Step 1: Retrieve top excerpts
+        retrieved_excerpts = llm_utils.retrieve_passages(
+            query=query,
+            doc_embeds=doc_embeds,
+            passages=passages,
+            processed_docs=processed_docs,
+            embed_model=embed_model,
+            max_docs=3,
+        )
+        progress(0.1, desc="Retrieved documents.")  # Update progress after retrieval
+
+        if not retrieved_excerpts:
+            print("No passages retrieved.")
+            progress(1, desc="No relevant passages found.")  # Update progress
+            return [
+                "### Document:\n-",
+                "<p><i>No relevant passages found.</i></p>",
+                "_",
+            ] * 3 + [
+                "### Summary\n_No passages found._"
+            ]  # Add summary placeholder
+
+        # Step 2: Process each excerpt individually
+        processed_data = []
+        num_excerpts = len(retrieved_excerpts)
+        for i, excerpt in enumerate(retrieved_excerpts):
+            # Update progress description before processing each excerpt
+            progress(
+                0.1 + (i * 0.8 / num_excerpts),
+                desc=f"Processing excerpt {i + 1}/{num_excerpts}...",
+            )
+            # Pass the globally initialized hf_client
+            processed_result = llm_utils.process_single_excerpt(
+                i, excerpt, query, hf_client
+            )
+            processed_data.append({"excerpt": excerpt, **processed_result})
+
+        progress(
+            0.9, desc="Formatting results..."
+        )  # Mark processing complete, start formatting
+
+        # Step 3: Format results for Gradio output components
+        final_outputs = [
+            "### Document:\n-",
+            "<p><i>No result</i></p>",
+            "_No passage text_",
+        ] * 3 + [
+            "### Summary\n_Generating..._"
+        ]  # Reset outputs + summary placeholder
+
+        global_quote_counter = 0  # Initialize global counter
+        snippets_for_llm_summary = []  # List to store formatted snippets for the LLM
+        all_spans_with_hover_info = (
+            {}
+        )  # Dict to store spans per passage index: {0: [(s, e, info), ...], 1: ...}
+
+        for i in range(min(len(processed_data), 3)):
+            result = processed_data[i]
+            excerpt = result["excerpt"]
+            citations = result["citations"]
+            parse_successful = result["parse_successful"]
+            raw_error_response = result["raw_error_response"]
+            passage_text = excerpt.get("passage_text", "")
+            doc_url = excerpt.get("document_url", "#")
+
+            # 1. Format Document URL Markdown
+            doc_url_md = (
+                f"### Document {i+1}:\n[{html.escape(doc_url)}]({html.escape(doc_url)})"
+            )
+
+            # 2. Format Quotes HTML
+            quotes_html_parts = []
+            if parse_successful and citations:
+                quotes_html_parts.append(
+                    "<p><strong>Relevant Quotes:</strong> (hover for details)</p>"
+                )
+                quotes_html_parts.append(
+                    "<ul style='list-style-type: none; margin-left: 10px; padding-left: 0;'>"  # Use none for list type
+                )
+                for cit in citations:
+                    global_quote_counter += 1  # Increment counter
+                    quote_id = global_quote_counter
+                    # Store id in citation dict (optional, but might be useful later)
+                    cit["global_id"] = quote_id
+
+                    quote = cit.get("quote", "N/A")
+                    context = cit.get("context", "N/A")
+                    relevance = cit.get("relevance", "N/A")
+                    hover_text = f"Context: {html.escape(context, quote=True)}\nRelevance: {html.escape(relevance, quote=True)}"
+                    # Update HTML to include the ID
+                    quotes_html_parts.append(
+                        f"<li style='margin-bottom: 5px;' title='{hover_text}'>[{quote_id}]: <i>{html.escape(quote)}</i></li>"
+                    )
+
+                    # Prepare hover text for the highlighted span in the passage
+                    span_hover_text = f"Quote ID: {quote_id}\nContext: {context}\nRelevance: {relevance}"
+                    # Get spans for this specific citation
+                    citation_spans = cit.get("char_spans", [])
+                    # Associate hover text with these spans for the current passage (index i)
+                    if i not in all_spans_with_hover_info:
+                        all_spans_with_hover_info[i] = []
+                    for start, end in citation_spans:
+                        all_spans_with_hover_info[i].append(
+                            (start, end, span_hover_text)
+                        )
+
+                    # Add formatted snippet to list if parsing was successful
+                    if (
+                        parse_successful
+                    ):  # Ensure we only add successfully parsed citations
+                        snippets_for_llm_summary.append(
+                            {
+                                "id": quote_id,
+                                "context": context,
+                                "relevance": relevance,
+                                "quote": quote,
+                                "document_url": doc_url,  # Added document URL here
+                            }
+                        )
+                quotes_html_parts.append("</ul>")
+            elif not parse_successful and raw_error_response:
+                quotes_html_parts.append(
+                    "<p style='color: red;'><strong>Error parsing citations:</strong></p>"
+                )
+                # Limit error display length
+                error_display = html.escape(raw_error_response[:1000]) + (
+                    "..." if len(raw_error_response) > 1000 else ""
+                )
+                quotes_html_parts.append(
+                    f"<pre style='background-color: #f8d7da; color: #721c24; padding: 5px; border-radius: 4px; white-space: pre-wrap; word-wrap: break-word;'><code>{error_display}</code></pre>"
+                )
+            else:
+                quotes_html_parts.append("<p><i>No specific quotes identified.</i></p>")
+            quotes_html = "".join(quotes_html_parts)
+            if not quotes_html:
+                quotes_html = "<p><i>No quotes processed.</i></p>"
+
+            # 3. Format Passage Markdown using the collected spans with hover info
+            spans_for_this_passage = all_spans_with_hover_info.get(
+                i, []
+            )  # Get spans for index i
+            passage_md = interface_utils.generate_highlighted_markdown(
+                passage_text, spans_for_this_passage
+            )
+            if not passage_md:
+                passage_md = "_Passage text unavailable._"
+
+            # Update the final_outputs list
+            final_outputs[i * 3 + 0] = doc_url_md
+            final_outputs[i * 3 + 1] = quotes_html
+            final_outputs[i * 3 + 2] = passage_md
+
+        # Step 4: Generate LLM summary
+        progress(0.95, desc="Generating summary...")  # New progress step
+        summary_text = "### Summary\n_Error generating summary._"  # Default error text
+        if snippets_for_llm_summary:
+            # Create a lookup for quote details by ID
+            snippet_lookup = {s["id"]: s for s in snippets_for_llm_summary}
+
+            # Pass the globally initialized hf_client
+            summary_result = llm_utils.generate_summary_answer(
+                snippets=snippets_for_llm_summary, query=query, hf_client=hf_client
+            )
+            if summary_result["parse_successful"]:
+                summary_items = []
+                for sentence_data in summary_result["answer_sentences"]:
+                    sentence = sentence_data.get("sentence", "")
+                    citation_ids = sentence_data.get("citations", [])
+
+                    # Generate HTML links for citations
+                    citation_links = []
+                    for c_id in citation_ids:
+                        # Look up snippet details
+                        snippet_info = snippet_lookup.get(c_id)
+                        if snippet_info:
+                            url = snippet_info.get("document_url", "#")
+                            quote_text = snippet_info.get("quote", "")
+                            escaped_quote = html.escape(quote_text, quote=True)
+                            link = f"<a href='{url}' title='{escaped_quote}' target='_blank'>[{c_id}]</a>"
+                            citation_links.append(link)
+                        else:
+                            citation_links.append(
+                                f"[{c_id}]"
+                            )  # Fallback if ID not found
+
+                    # Format citation string like [link1, link2]
+                    citation_str = (
+                        f" [{', '.join(citation_links)}]" if citation_links else ""
+                    )
+                    if sentence:
+                        # Append sentence and linked citations
+                        summary_items.append(f"{html.escape(sentence)}{citation_str}")
+                if summary_items:
+                    summary_text = "### Generated Answer\n" + " ".join(summary_items)
+                else:
+                    summary_text = (
+                        "### Generated Answer\n_LLM did not generate any sentences._"
+                    )
+            else:
+                # Display parsing error from summary LLM
+                summary_text = (
+                    "### Summary Generation Error\n"
+                    + f"<p style='color: red;'>Could not generate summary:</p><pre style='background-color: #f8d7da; color: #721c24; padding: 5px; border-radius: 4px; white-space: pre-wrap; word-wrap: break-word;'><code>{html.escape(summary_result['raw_error_response'])}</code></pre>"
+                )
+        else:
+            summary_text = "### Summary\n_No valid quotes found to generate summary._"
+
+        final_outputs[9] = summary_text  # Add summary to the 10th slot
+
+        progress(1, desc="Done!")  # Final progress update
+        return final_outputs
+
+    except Exception as e:
+        print(f"Error in get_results: {e}")
+        # Display error in the first result slot, update summary slot
+        error_outputs = (
+            [
+                "### Error Processing Query",
+                f"<p style='color: red;'>An unexpected error occurred: {html.escape(str(e))}</p>",
+                "_Error details above._",
+            ]
+            + ["-"] * 6
+            + ["### Summary\n_Error occurred._"]
+        )  # Add summary placeholder for error
+        progress(1, desc="Error!")
+        return error_outputs
+
+
+# --- Gradio Interface ---
+
+# Define custom CSS for scrollable accordion content
+custom_css = """
+.scrollable-passage-content { 
+    max-height: 25vh;  /* Or a fixed height like 200px */
+    overflow-y: auto !important; /* Add !important to override potential conflicts */
+    display: block; /* Ensure it behaves like a block element */
+    padding: 10px; /* Add some padding */
+    background-color: #f9f9f9; /* Match previous background */
+    border: 1px solid #ddd; /* Match previous border */
+    border-radius: 5px; /* Match previous radius */
+    white-space: pre-wrap; /* Preserve whitespace */
+    word-wrap: break-word; /* Wrap long words */
+}
+"""
+
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft(), title="🤗 Policy Docs QA") as demo:
+    gr.Markdown("# 🤗 Policy Docs QA")
+    gr.Markdown(
+        "Ask a question about the loaded policy documents to retrieve relevant passages and quotes."
+    )  # Added description
+
+    with gr.Row():
+        # Column 1: Input (scale=1)
+        with gr.Column(scale=1):
+            question_input = gr.Textbox(
+                label="Question",
+                placeholder="Enter your question...",
+                lines=5,
+            )
+            # Add example questions
+            example_questions = [
+                "What role does replicable evaluation play in AI regulation?",
+                "How does dataset transparency help address risks of accidents and abuse of AI systems?",
+            ]
+            gr.Examples(
+                examples=example_questions,
+                inputs=question_input,
+                label="Example Questions",  # Optional label
+            )
+            submit_button = gr.Button("Get Answer")
+
+        # Column 2: Results (scale=3)
+        with gr.Column(scale=3):
+            gr.Markdown("## Retrieved Passages")  # Changed heading
+
+            result_outputs = []  # Rename list for clarity
+            for i in range(3):  # Create 3 result sections
+                with gr.Group():
+                    gr.Markdown(f"**Result {i+1}**")
+                    doc_url_md = gr.Markdown(
+                        value="### Document:\n-", label=f"Document URL {i+1}"
+                    )
+                    quotes_html = gr.HTML(
+                        value="<p><i>Quotes will appear here...</i></p>",
+                        label=f"Quotes {i+1}",
+                    )
+                    with gr.Accordion(f"Full Passage Context {i+1}", open=False):
+                        passage_md = gr.Markdown(
+                            value="_Passage text will appear here..._",
+                            label=f"Passage {i+1}",
+                            elem_classes="scrollable-passage-content",  # Assign the class
+                        )
+                    result_outputs.extend([doc_url_md, quotes_html, passage_md])
+
+        # Column 3: Summary (scale=2)
+        with gr.Column(scale=2):
+            gr.Markdown("## Summary")
+            gr.Markdown("***⚠️ Warning:** The text below is generated by an LLM and might not accurately reflect the policy documents. Citation links are determined by the same LLM and provided for conveninece only. For reliable information, go directly to the Retrieved Passages left.*")
+            summary_output = gr.Markdown(
+                value="_Summary will appear here..._", label="Summary"
+            )
+
+    # Combine all output components for the click function
+    all_outputs = result_outputs + [summary_output]
+    assert (
+        len(all_outputs) == 10
+    ), "Incorrect number of total output components created."
+
+    submit_button.click(
+        fn=get_results,
+        inputs=question_input,
+        outputs=all_outputs,  # Pass the list of 10 components
+    )
+
+# --- Launch App ---
+if __name__ == "__main__":
+    demo.launch()  # share=True for public link