Upload app.py

app.py

@@ -8,9 +8,10 @@ import anthropic
 import google.generativeai as genai
 import gradio as gr
 from gliner import GLiNER
-import traceback 
-from collections import defaultdict, Counter # Import Counter for counting
+import traceback
+from collections import defaultdict, Counter
 import re
+import os # Make sure this import is at the top of your file
 
 # 🧠 Supported models and their providers
 MODEL_OPTIONS = {
@@ -72,9 +73,9 @@ def generate_from_prompt(prompt, provider, key_dict):
     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", 
+    "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"
 ]
 
@@ -83,6 +84,22 @@ MAX_CATEGORIES = 8
 with gr.Blocks(title="Historical Text Analysis Tool", css=".prose { word-break: break-word; }") as demo:
     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("---")
+    
     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():
@@ -99,21 +116,26 @@ with gr.Blocks(title="Historical Text Analysis Tool", css=".prose { word-break:
     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. Uncheck any you do not want, or use the 'Deselect All' button for that category.")
+    gr.Markdown("Click on a category to see its keywords. Use the buttons to select or deselect all keywords for that category.")
     
-    dynamic_components = []
+    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)
-                dynamic_components.append((acc, cg, deselect_btn))
+                category_components.append((acc, cg, select_btn, deselect_btn))
     
     gr.Markdown("### 2. Include Standard Keywords (Optional)")
     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.")
-        deselect_ner_btn = gr.Button("Deselect All", size="sm")
+        # --- 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)")
     with gr.Group():
@@ -133,7 +155,6 @@ with gr.Blocks(title="Historical Text Analysis Tool", css=".prose { word-break:
             debug_output = gr.Textbox(label="Extraction Log", interactive=False, lines=8)
     
     # --- Backend Functions ---
-    import os # Make sure this import is at the top of your file
 
     def handle_generate(topic, provider, openai_k, anthropic_k, google_k):
         # This function provides instant "working..." feedback
@@ -170,24 +191,25 @@ with gr.Blocks(title="Historical Text Analysis Tool", css=".prose { word-break:
             updates = {}
             categories = list(framework.items())
             for i in range(MAX_CATEGORIES):
-                accordion_comp, checkbox_comp, button_comp = dynamic_components[i]
+                accordion_comp, checkbox_comp, sel_btn, desel_btn = 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[button_comp] = gr.update(visible=True)
+                    updates[sel_btn] = gr.update(visible=True)
+                    updates[desel_btn] = gr.update(visible=True)
                 else:
                     updates[accordion_comp] = gr.update(visible=False)
                     updates[checkbox_comp] = gr.update(visible=False)
-                    updates[button_comp] = gr.update(visible=False)
+                    updates[sel_btn] = gr.update(visible=False)
+                    updates[desel_btn] = 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))
 
-    # --- THIS IS THE UPDATED FUNCTION ---
     def match_entities(text, ner_labels, custom_label_text, threshold, *selected_keywords):
         debug_info = []
         if gliner_model is None:
@@ -221,19 +243,16 @@ with gr.Blocks(title="Historical Text Analysis Tool", css=".prose { word-break:
 
         highlighted_entities = [{"start": ent["start"], "end": ent["end"], "entity": ent["label"]} for ent in unique_entities]
         
-        # --- NEW LOGIC FOR AGGREGATED, TABLE-BASED RESULTS ---
-        # 1. Count occurrences of each unique phrase (case-insensitively)
         aggregated_matches = defaultdict(Counter)
-        original_casing_map = {} # To store the original casing of the first instance of a phrase
+        original_casing_map = {} 
 
         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) # Store original casing
+            original_casing_map.setdefault(match_text_lower, match_text) 
             
-        # 2. Build the new Markdown string with tables
         markdown_string = ""
         for label, counter in sorted(aggregated_matches.items()):
             total_matches = sum(counter.values())
@@ -242,7 +261,6 @@ with gr.Blocks(title="Historical Text Analysis Tool", css=".prose { word-break:
             markdown_string += "| Found Phrase | Occurrences |\n"
             markdown_string += "|--------------|-------------|\n"
             
-            # Sort phrases by most frequent first
             for phrase_lower, count in counter.most_common():
                 original_phrase = original_casing_map[phrase_lower]
                 markdown_string += f"| {original_phrase} | {count} |\n"
@@ -254,22 +272,42 @@ with gr.Blocks(title="Historical Text Analysis Tool", css=".prose { word-break:
         return {"text": text, "entities": highlighted_entities}, markdown_string, "\n".join(debug_info)
 
     # --- Wire up UI events ---
-    generate_btn.click(
-        fn=handle_generate,
-        inputs=[topic, provider, openai_key, anthropic_key, google_key],
-        outputs=[generate_btn] + [comp for pair in dynamic_components for comp in pair]
-    )
     
+    # 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"
+    }
+    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)
+
+    # --- 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])
-    for _, cg, btn in dynamic_components:
-        btn.click(fn=deselect_all, inputs=None, outputs=[cg])
 
+    # --- 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
     match_btn.click(
         fn=match_entities,
-        inputs=[text_input, ner_output, custom_labels, threshold_slider] + [cg for acc, cg, btn in dynamic_components],
-        outputs=[matched_output, detailed_results_output, debug_output]
+        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"
     )
 
 demo.launch(share=True, debug=True)