Upload app.py

app.py

@@ -1,6 +1,6 @@
 # 📚 Install dependencies
 # Make sure to run this in your environment if you haven't already
-# !pip install openai anthropic google-generativeai gradio transformers torch gliner --quiet
+# !pip install openai anthropic google-generativeai gradio transformers torch gliner pandas --quiet
 
 # ⚙️ Imports
 import openai
@@ -11,7 +11,9 @@ from gliner import GLiNER
 import traceback
 from collections import defaultdict, Counter
 import re
-import os # Make sure this import is at the top of your file
+import os
+import pandas as pd
+import tempfile
 
 # 🧠 Supported models and their providers
 MODEL_OPTIONS = {
@@ -32,206 +34,159 @@ except Exception as e:
     print(f"FATAL ERROR: Could not load GLiNER model. The app will not be able to find entities. Error: {e}")
     gliner_model = None
 
-# 🧠 Prompt for generating the research framework
-HIERARCHICAL_PROMPT_TEMPLATE = """
-You are a helpful research assistant. For the historical topic: **"{topic}"**, your job is to suggest a research framework.
-
-**Instructions:**
-1.  First, think of 4-6 **Conceptual Categories** that are useful for analyzing this topic (e.g., 'Forms of Protest', 'Key Demands').
-2.  For each category, list the specific **Keywords** someone could search for in a text.
-3.  **Crucial Rule for Keywords:** Use the most basic, fundamental form (e.g., `Petition`, not `Political Petition`).
-
-**Output Format:**
-Use Markdown. Each category must be a Level 3 Header (###), followed by a comma-separated list of its keywords.
-
-### Example Category 1
-- Keyword A, Keyword B, Keyword C
-### Example Category 2
-- Keyword D, Keyword E
-"""
-
-# 🧠 Generator Function
-def generate_from_prompt(prompt, provider, key_dict):
-    provider_id = MODEL_OPTIONS.get(provider)
-    api_key = key_dict.get(f"{provider_id}_key")
-    if not api_key:
-        raise ValueError(f"API key for {provider} not found.")
-
-    if provider_id == "openai":
-        client = openai.OpenAI(api_key=api_key)
-        response = client.chat.completions.create(model="gpt-4o", messages=[{"role": "user", "content": prompt}], temperature=0.2)
-        return response.choices[0].message.content.strip()
-    elif provider_id == "anthropic":
-        client = anthropic.Anthropic(api_key=api_key)
-        response = client.messages.create(model="claude-3-opus-20240229", max_tokens=1024, messages=[{"role": "user", "content": prompt}])
-        return response.content[0].text.strip()
-    elif provider_id == "google":
-        genai.configure(api_key=api_key)
-        model = genai.GenerativeModel('gemini-1.5-pro-latest')
-        response = model.generate_content(prompt)
-        return response.text.strip()
-    return ""
-
-TRADITIONAL_NER_LABELS = [
-    "Person", "Organisation", "Country / City / State", "Location",
-    "Nationality or Group", "Date", "Event", "Law / Legal Document",
-    "Product", "Facility", "Work of Art", "Language", "Time", "Percentage",
-    "Money / Currency", "Quantity / Measurement", "Ordinal Number", "Cardinal Number"
-]
-
+# --- Prompt and other constants remain the same ---
+HIERARCHICAL_PROMPT_TEMPLATE = "..." # (Keeping this collapsed for brevity, no changes needed)
+TRADITIONAL_NER_LABELS = ["..."] # (Keeping this collapsed for brevity, no changes needed)
 MAX_CATEGORIES = 8
 
 with gr.Blocks(title="Historical Text Analysis Tool", css=".prose { word-break: break-word; }") as demo:
+    # --- UI remains the same up to the output tabs ---
     gr.Markdown("# Historical Text Analysis Tool")
-    
-    # --- NEW: Added introductory text ---
-    gr.Markdown(
-        """
-        **Welcome! This tool uses two different kinds of AI to help you quickly analyze documents.**
-
-        1.  **The "Creative Assistant" (Step 1: OpenAI, Anthropic, Google):**
-            When you enter a topic, this AI acts like a research assistant. It brainstorms and **suggests** useful categories and keywords for your analysis. It's the idea generator.
-
-        2.  **The "Expert Searcher" (Step 2: GLiNER):**
-            After you've chosen your keywords, this specialized AI meticulously **finds** every single match in the text you provide. It's a fast and precise search tool that runs locally.
-
-        **Pro Tip:** After the analysis, you can manually add or correct a label! In the "Highlighted Text" tab, just click on any word or phrase, type your new label, and press Enter.
-        """
-    )
+    gr.Markdown("...") # Welcome text collapsed for brevity
     gr.Markdown("---")
-    
     gr.Markdown("## Step 1: Get Keyword Ideas")
-    gr.Markdown("Start by entering a topic. The AI will populate a research framework with suggested categories and keywords to guide your analysis.")
     with gr.Row():
-        topic = gr.Textbox(label="Enter Historical Topic", placeholder="e.g., The Chartist Movement, The Protestant Reformation")
+        topic = gr.Textbox(label="Enter Historical Topic", placeholder="e.g., The Chartist Movement")
         provider = gr.Dropdown(choices=list(MODEL_OPTIONS.keys()), label="Choose AI Model")
     with gr.Row():
-        openai_key = gr.Textbox(label="OpenAI API Key", type="password", placeholder="Required for OpenAI")
-        anthropic_key = gr.Textbox(label="Anthropic API Key", type="password", placeholder="Required for Anthropic")
-        google_key = gr.Textbox(label="Google API Key", type="password", placeholder="Required for Google")
-    
+        openai_key = gr.Textbox(label="OpenAI API Key", type="password")
+        anthropic_key = gr.Textbox(label="Anthropic API Key", type="password")
+        google_key = gr.Textbox(label="Google API Key", type="password")
     generate_btn = gr.Button("Suggest Categories and Keywords", variant="primary")
-
-    gr.Markdown("--- \n## Step 2: Build Your Search and Analyze Text")
-    gr.Markdown("The AI's suggestions will appear below. Build your final list of keywords, then paste your text to find all the matches.")
-    
-    gr.Markdown("### 1. Review AI-Suggested Keywords")
-    gr.Markdown("Click on a category to see its keywords. Use the buttons to select or deselect all keywords for that category.")
     
+    gr.Markdown("--- \n## Step 2: Build Your Search and Analyze Text")
     category_components = []
     with gr.Column():
         for i in range(MAX_CATEGORIES):
             with gr.Accordion(f"Category {i+1}", visible=False) as acc:
                 with gr.Row():
                     cg = gr.CheckboxGroup(label="Keywords", interactive=True, container=False, scale=4)
-                    # --- NEW: Added Select All button for categories ---
-                    select_btn = gr.Button("Select All", size="sm", scale=1, min_width=80)
-                    deselect_btn = gr.Button("Deselect All", size="sm", scale=1, min_width=80)
-                category_components.append((acc, cg, select_btn, deselect_btn))
-    
-    gr.Markdown("### 2. Include Standard Keywords (Optional)")
+                    toggle_btn = gr.Button("Deselect All", size="sm", scale=1, min_width=100)
+                category_components.append((acc, cg, toggle_btn))
     with gr.Group():
-        ner_output = gr.CheckboxGroup(choices=TRADITIONAL_NER_LABELS, value=TRADITIONAL_NER_LABELS, label="Standard Search Terms", info="Common categories like people, places, and specific organizations.")
-        # --- NEW: Added Select All button for standard keywords ---
-        with gr.Row():
-            select_ner_btn = gr.Button("Select All", size="sm")
-            deselect_ner_btn = gr.Button("Deselect All", size="sm")
-
-    gr.Markdown("### 3. Add Your Own Keywords (Optional)")
+        ner_output = gr.CheckboxGroup(choices=TRADITIONAL_NER_LABELS, value=TRADITIONAL_NER_LABELS, label="Standard Search Terms")
+        toggle_ner_btn = gr.Button("Deselect All", size="sm")
     with gr.Group():
-        gr.Markdown("**Add any other keywords**")
-        custom_labels = gr.Textbox(label=None, placeholder="e.g., Technology, Weapon, Secret Society... (separated by commas)", show_label=False)
-    
-    threshold_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.4, step=0.05, label="Confidence Threshold", info="This controls how strict the search is. Lower to find more matches (less strict). Raise for fewer, more precise matches (more strict).")
-    text_input = gr.Textbox(label="Paste Your Full Text Here for Analysis", lines=10, placeholder="Paste a historical document, an article, or a chapter...")
+        custom_labels = gr.Textbox(label="Add Your Own Keywords (Optional)", placeholder="e.g., Technology, Weapon... (separated by commas)")
+    threshold_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.4, step=0.05, label="Confidence Threshold", info="Controls how 'sure' the AI needs to be. Lower finds more potential matches, higher finds only the most certain ones.")
+    text_input = gr.Textbox(label="Paste Your Full Text Here for Analysis", lines=10)
     match_btn = gr.Button("Find Keywords in Text", variant="primary")
     
+    # --- NEW: Add state variables to hold data between function calls ---
+    # This holds the original text for updates
+    text_state = gr.State()
+    # This holds the results DataFrame for updates and downloads
+    dataframe_state = gr.State()
+
     with gr.Tabs():
         with gr.TabItem("Highlighted Text"):
-            matched_output = gr.HighlightedText(label="Keyword Matches", interactive=True)
+            matched_output = gr.HighlightedText(
+                label="Keyword Matches",
+                interactive=True,
+                show_legend=True
+            )
         with gr.TabItem("Detailed Results"):
-            detailed_results_output = gr.Markdown(label="List of Matches per Keyword")
+            # --- CHANGE: Using gr.DataFrame for a clean table output ---
+            detailed_results_output = gr.DataFrame(
+                headers=["Category", "Found Phrase", "Occurrences"],
+                datatype=["str", "str", "number"],
+                wrap=True,
+                label="Aggregated Results"
+            )
+            # --- NEW: Download button and hidden file component ---
+            download_button = gr.Button("Download Results as CSV", visible=False)
+            download_file = gr.File(label="Download", visible=False)
+
         with gr.TabItem("Debug Info"):
             debug_output = gr.Textbox(label="Extraction Log", interactive=False, lines=8)
     
     # --- Backend Functions ---
-
     def handle_generate(topic, provider, openai_k, anthropic_k, google_k):
-        # This function provides instant "working..." feedback
-        yield {
-            generate_btn: gr.update(value="Generating...", interactive=False)
-        }
-        
+        # ... (This function remains unchanged) ...
+        yield {generate_btn: gr.update(value="Consulting the Archives...", interactive=False)}
         try:
-            # On Hugging Face, use secure secrets. Locally, use the text boxes.
-            key_dict = {
-                "openai_key": os.environ.get("OPENAI_API_KEY", openai_k),
-                "anthropic_key": os.environ.get("ANTHROPIC_API_KEY", anthropic_k),
-                "google_key": os.environ.get("GOOGLE_API_KEY", google_k)
-            }
-
+            key_dict = {"openai_key": os.environ.get("OPENAI_API_KEY", openai_k), "anthropic_key": os.environ.get("ANTHROPIC_API_KEY", anthropic_k), "google_key": os.environ.get("GOOGLE_API_KEY", google_k)}
             provider_id = MODEL_OPTIONS.get(provider)
-            if not topic or not provider or not key_dict.get(f"{provider_id}_key"):
-                raise gr.Error("Topic, Provider, and the correct API Key are required.")
-            
+            if not topic or not provider or not key_dict.get(f"{provider_id}_key"): raise gr.Error("Topic, Provider, and the correct API Key are required.")
             prompt = HIERARCHICAL_PROMPT_TEMPLATE.format(topic=topic)
             raw_framework = generate_from_prompt(prompt, provider, key_dict)
             framework = defaultdict(list)
             current_category = None
             for line in raw_framework.split('\n'):
                 line = line.strip()
-                if line.startswith("###"):
-                    current_category = line.replace("###", "").strip()
-                elif line.startswith("-") and current_category:
-                    entities = line.replace("-", "").strip()
-                    framework[current_category].extend([e.strip() for e in entities.split(',') if e.strip()])
-            if not framework:
-                raise gr.Error("AI failed to generate categories. Please try again.")
-
+                if line.startswith("###"): current_category = line.replace("###", "").strip()
+                elif line.startswith("-") and current_category: framework[current_category].extend([e.strip() for e in line.replace("-", "").strip().split(',') if e.strip()])
+            if not framework: raise gr.Error("AI failed to generate categories. Please try again.")
             updates = {}
             categories = list(framework.items())
             for i in range(MAX_CATEGORIES):
-                accordion_comp, checkbox_comp, sel_btn, desel_btn = category_components[i]
+                accordion_comp, checkbox_comp, toggle_btn_comp = category_components[i]
                 if i < len(categories):
                     category, entities = categories[i]
                     sorted_entities = sorted(list(set(entities)))
                     updates[accordion_comp] = gr.update(label=category, visible=True)
                     updates[checkbox_comp] = gr.update(choices=sorted_entities, value=sorted_entities, visible=True)
-                    updates[sel_btn] = gr.update(visible=True)
-                    updates[desel_btn] = gr.update(visible=True)
+                    updates[toggle_btn_comp] = gr.update(visible=True, value="Deselect All")
                 else:
                     updates[accordion_comp] = gr.update(visible=False)
                     updates[checkbox_comp] = gr.update(visible=False)
-                    updates[sel_btn] = gr.update(visible=False)
-                    updates[desel_btn] = gr.update(visible=False)
+                    updates[toggle_btn_comp] = gr.update(visible=False)
             updates[generate_btn] = gr.update(value="Suggest Categories and Keywords", interactive=True)
             yield updates
         except Exception as e:
             yield {generate_btn: gr.update(value="Suggest Categories and Keywords", interactive=True)}
             raise gr.Error(str(e))
 
-    def match_entities(text, ner_labels, custom_label_text, threshold, *selected_keywords):
-        debug_info = []
-        if gliner_model is None:
-            raise gr.Error("GLiNER model failed to load at startup. Cannot analyze text. Please check the logs and restart the application.")
+    # --- NEW: Helper function to process entities into a DataFrame ---
+    def process_entities_to_df(entities, original_text):
+        """Takes a list of entities and the original text, and returns a pandas DataFrame."""
+        if not entities:
+            return pd.DataFrame(columns=["Category", "Found Phrase", "Occurrences"])
+        
+        # Extract text for each entity
+        found_phrases = []
+        for ent in entities:
+            found_phrases.append({
+                "Category": ent['entity'],
+                "Found Phrase": original_text[ent['start']:ent['end']]
+            })
+        
+        if not found_phrases:
+            return pd.DataFrame(columns=["Category", "Found Phrase", "Occurrences"])
+        
+        # Aggregate using pandas
+        df = pd.DataFrame(found_phrases)
+        aggregated_df = df.groupby(["Category", "Found Phrase"]).size().reset_index(name="Occurrences")
+        aggregated_df = aggregated_df.sort_values(by=["Category", "Occurrences"], ascending=[True, False])
+        
+        return aggregated_df
 
+    # --- UPDATED: `match_entities` now uses pandas and updates state ---
+    def match_entities(text, ner_labels, custom_label_text, threshold, *selected_keywords, progress=gr.Progress(track_tqdm=True)):
+        yield {
+            match_btn: gr.update(value="Searching...", interactive=False),
+            detailed_results_output: None,
+            download_button: gr.update(visible=False),
+            download_file: gr.update(visible=False)
+        }
+        if gliner_model is None: raise gr.Error("GLiNER model failed to load.")
+        
         labels_to_use = set()
+        if ner_labels: labels_to_use.update(ner_labels)
         for group in selected_keywords:
             if group: labels_to_use.update(group)
-        if ner_labels: labels_to_use.update(ner_labels)
         custom = {l.strip() for l in custom_label_text.split(',') if l.strip()}
         if custom: labels_to_use.update(custom)
-        
         final_labels = sorted(list(labels_to_use))
-        debug_info.append(f"🧠 Searching for {len(final_labels)} unique keywords.")
-        debug_info.append(f"⚙️ Confidence Threshold: {threshold}")
-
+        debug_info = [f"🧠 Searching for {len(final_labels)} unique keywords.", f"⚙️ Confidence Threshold: {threshold}"]
+        
         if not text or not final_labels:
-            return {"text": text, "entities": []}, "Please provide text and select keywords.", "\n".join(debug_info)
+            yield {match_btn: gr.update(value="Find Keywords in Text", interactive=True)}
+            return
         
         all_entities = []
         chunk_size, overlap = 1000, 50
-        for i in range(0, len(text), chunk_size - overlap):
+        for i in progress.tqdm(range(0, len(text), chunk_size - overlap), desc="Scanning Text..."):
             chunk = text[i : i + chunk_size]
             chunk_entities = gliner_model.predict_entities(chunk, final_labels, threshold=threshold)
             for ent in chunk_entities:
@@ -240,74 +195,83 @@ with gr.Blocks(title="Historical Text Analysis Tool", css=".prose { word-break:
         
         unique_entities = [dict(t) for t in {tuple(d.items()) for d in all_entities}]
         debug_info.append(f"📊 Found {len(unique_entities)} unique matches.")
-
+        
         highlighted_entities = [{"start": ent["start"], "end": ent["end"], "entity": ent["label"]} for ent in unique_entities]
         
-        aggregated_matches = defaultdict(Counter)
-        original_casing_map = {} 
+        # --- NEW: Use helper to create DataFrame ---
+        results_df = process_entities_to_df(highlighted_entities, text)
 
-        for ent in unique_entities:
-            match_text = text[ent['start']:ent['end']]
-            match_text_lower = match_text.lower()
-            
-            aggregated_matches[ent['label']][match_text_lower] += 1
-            original_casing_map.setdefault(match_text_lower, match_text) 
-            
-        markdown_string = ""
-        for label, counter in sorted(aggregated_matches.items()):
-            total_matches = sum(counter.values())
-            unique_phrases = len(counter)
-            markdown_string += f"### {label} (Total: {total_matches} | Unique: {unique_phrases})\n"
-            markdown_string += "| Found Phrase | Occurrences |\n"
-            markdown_string += "|--------------|-------------|\n"
-            
-            for phrase_lower, count in counter.most_common():
-                original_phrase = original_casing_map[phrase_lower]
-                markdown_string += f"| {original_phrase} | {count} |\n"
-            markdown_string += "\n"
-            
-        if not markdown_string:
-            markdown_string = "No keywords found. Try lowering the confidence threshold or changing keywords."
-
-        return {"text": text, "entities": highlighted_entities}, markdown_string, "\n".join(debug_info)
+        yield {
+            match_btn: gr.update(value="Find Keywords in Text", interactive=True),
+            matched_output: {"text": text, "entities": highlighted_entities},
+            detailed_results_output: results_df,
+            debug_output: "\n".join(debug_info),
+            download_button: gr.update(visible=True if not results_df.empty else False),
+            text_state: text, # Store original text in state
+            dataframe_state: results_df # Store dataframe in state
+        }
 
-    # --- Wire up UI events ---
+    # --- NEW: Function to update results when highlighted text is edited ---
+    def update_detailed_results(new_highlighted_entities, original_text):
+        """
+        This function is triggered when the user edits the HighlightedText component.
+        It re-calculates the DataFrame and updates the UI.
+        """
+        # new_highlighted_entities is the full value of the component, not just a diff
+        results_df = process_entities_to_df(new_highlighted_entities, original_text)
+        
+        return {
+            detailed_results_output: results_df,
+            dataframe_state: results_df, # Update the state for the download button
+            download_button: gr.update(visible=True if not results_df.empty else False),
+        }
+        
+    # --- NEW: Function to handle the file download ---
+    def download_results_as_csv(df):
+        """Saves the DataFrame to a temporary CSV file and returns its path."""
+        with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.csv', encoding='utf-8') as tmp:
+            df.to_csv(tmp.name, index=False)
+            return gr.update(value=tmp.name, visible=True)
+
+    # --- Event Wiring ---
+    def handle_toggle_click(button_text, all_choices):
+        if button_text == "Select All": return gr.update(value=all_choices), gr.update(value="Deselect All")
+        else: return gr.update(value=[]), gr.update(value="Select All")
+    def update_button_on_check(selections):
+        return gr.update(value="Select All") if not selections else gr.update(value="Deselect All")
     
-    # NEW: Handle "Enter" key press on the topic textbox and show progress bar
-    submit_event_args = {
-        "fn": handle_generate,
-        "inputs": [topic, provider, openai_key, anthropic_key, google_key],
-        "outputs": [generate_btn] + [comp for pair in category_components for comp in pair],
-        "show_progress": "full"
-    }
+    submit_event_args = {"fn": handle_generate, "inputs": [topic, provider, openai_key, anthropic_key, google_key], "outputs": [generate_btn] + [comp for pair in category_components for comp in pair]}
     generate_btn.click(**submit_event_args)
     topic.submit(**submit_event_args)
-    
-    # --- NEW: Helper functions for select/deselect ---
-    def deselect_all():
-        return gr.update(value=[])
 
-    def select_all_ner():
-        return gr.update(value=TRADITIONAL_NER_LABELS)
-        
-    def select_all_from_group(checkbox_group_state):
-        return gr.update(value=checkbox_group_state.choices)
+    toggle_ner_btn.click(fn=handle_toggle_click, inputs=[toggle_ner_btn, gr.State(TRADITIONAL_NER_LABELS)], outputs=[ner_output, toggle_ner_btn])
+    ner_output.change(fn=update_button_on_check, inputs=[ner_output], outputs=[toggle_ner_btn])
 
-    # --- NEW: Wire up select/deselect for standard keywords ---
-    select_ner_btn.click(fn=select_all_ner, inputs=None, outputs=[ner_output])
-    deselect_ner_btn.click(fn=deselect_all, inputs=None, outputs=[ner_output])
-
-    # --- UPDATED: Wire up select/deselect for dynamic categories ---
-    for acc, cg, select_btn, deselect_btn in category_components:
-        select_btn.click(fn=select_all_from_group, inputs=[cg], outputs=[cg])
-        deselect_btn.click(fn=deselect_all, inputs=None, outputs=[cg])
-
-    # NEW: Show progress bar for the matching process
+    def create_toggle_handler(cg_component):
+        def handler(button_text): return handle_toggle_click(button_text, cg_component.choices)
+        return handler
+    for acc, cg, toggle_btn in category_components:
+        toggle_btn.click(fn=create_toggle_handler(cg), inputs=[toggle_btn], outputs=[cg, toggle_btn])
+        cg.change(fn=update_button_on_check, inputs=[cg], outputs=[toggle_btn])
+    
     match_btn.click(
         fn=match_entities,
-        inputs=[text_input, ner_output, custom_labels, threshold_slider] + [cg for acc, cg, sel, desel in category_components],
-        outputs=[matched_output, detailed_results_output, debug_output],
-        show_progress="full"
+        inputs=[text_input, ner_output, custom_labels, threshold_slider] + [cg for acc, cg, btn in category_components],
+        # --- CHANGE: Added new state and download components to outputs ---
+        outputs=[match_btn, matched_output, detailed_results_output, debug_output, download_button, download_file, text_state, dataframe_state]
+    )
+
+    # --- NEW: Wire up the dynamic update and download events ---
+    matched_output.change(
+        fn=update_detailed_results,
+        inputs=[matched_output, text_state],
+        outputs=[detailed_results_output, dataframe_state, download_button]
+    )
+    
+    download_button.click(
+        fn=download_results_as_csv,
+        inputs=[dataframe_state],
+        outputs=[download_file]
     )
 
 demo.launch(share=True, debug=True)