Update app.py

app.py

@@ -1,13 +1,13 @@
 import gradio as gr
 from transformers import pipeline
 
-# -------------------------------
+# --------------------------------------
 # MODEL CONFIGURATION
-# -------------------------------
+# --------------------------------------
 
 model_name = "ELHACHYMI/bert-ner"
 
-example_sent = "Bill Gates founded Microsoft in the United States."
+default_example = "Bill Gates founded Microsoft in the United States."
 
 example_list = [
     "Elon Musk is the CEO of SpaceX and Tesla.",
@@ -17,103 +17,166 @@ example_list = [
     "Apple Inc. was created by Steve Jobs in Cupertino.",
 ]
 
-# Model info for usage display
+# Correct mapping from LABEL_X to entity names
+ENTITY_MAP = {
+    "LABEL_0": "O",
+    "LABEL_1": "PER",
+    "LABEL_2": "PER",
+    "LABEL_3": "ORG",
+    "LABEL_4": "ORG",
+    "LABEL_5": "LOC",
+    "LABEL_6": "LOC",
+    "LABEL_7": "MISC",
+    "LABEL_8": "MISC",
+}
+
+# Colors for HTML rendering
+COLOR_MAP = {
+    "PER": "#ff6666",   # Red
+    "ORG": "#66b3ff",   # Blue
+    "LOC": "#99ff99",   # Green
+    "MISC": "#ffcc66",  # Orange
+}
+
 model_info = {
     model_name: {
         "link": f"https://huggingface.co/{model_name}",
         "usage": f"""from transformers import pipeline
 
 ner = pipeline("ner", model="{model_name}", aggregation_strategy="simple")
-text = "{example_sent}"
-print(ner(text))""",
+text = "{default_example}"
+print(ner(text))"""
     }
 }
 
-# Load model
 ner_pipeline = pipeline("ner", model=model_name, aggregation_strategy="simple")
 
 
-# -------------------------------
+# --------------------------------------
 # DISPLAY MODEL INFO
-# -------------------------------
+# --------------------------------------
 
 def display_model_info(model_name):
     info = model_info[model_name]
     usage_code = info["usage"]
-    link_button = f'[ Open model page: {model_name}]({info["link"]})'
-    return usage_code, link_button
-
+    model_link = f'[Open model page: {model_name}]({info["link"]})'
+    return usage_code, model_link
 
-# -------------------------------
-# NER ANALYSIS FUNCTION
-# -------------------------------
 
-def analyze_text(text):
-    ner_results = ner_pipeline(text)
+# --------------------------------------
+# HIGHLIGHTED MODE FUNCTION
+# --------------------------------------
 
-    highlighted_text = []
+def analyze_highlight(text):
+    results = ner_pipeline(text)
+    highlighted = []
     last_idx = 0
 
-    for entity in ner_results:
-        start = entity["start"]
-        end = entity["end"]
-        label = entity["entity_group"]
+    for ent in results:
+        start, end = ent["start"], ent["end"]
+        raw_label = ent["entity_group"]
+        label = ENTITY_MAP.get(raw_label, raw_label)
 
-        # Text before entity
         if start > last_idx:
-            highlighted_text.append((text[last_idx:start], None))
+            highlighted.append((text[last_idx:start], None))
 
-        # Entity text
-        highlighted_text.append((text[start:end], label))
+        highlighted.append((text[start:end], label))
         last_idx = end
 
-    # Remaining text
     if last_idx < len(text):
-        highlighted_text.append((text[last_idx:], None))
+        highlighted.append((text[last_idx:], None))
+
+    return highlighted
+
+
+# --------------------------------------
+# HTML COLOR MODE FUNCTION
+# --------------------------------------
+
+def colorize_text(text):
+    results = ner_pipeline(text)
+    html = ""
+    last = 0
+
+    for ent in results:
+        start, end = ent["start"], ent["end"]
+        raw_label = ent["entity_group"]
+        label = ENTITY_MAP.get(raw_label, raw_label)
 
-    return highlighted_text
+        if start > last:
+            html += text[last:start]
 
+        color = COLOR_MAP.get(label, "#dddddd")
+        html += (
+            f"<span style='background-color:{color}; padding:2px; border-radius:3px;'>"
+            f"{text[start:end]} ({label})"
+            f"</span>"
+        )
 
-# -------------------------------
-# GRADIO INTERFACE
-# -------------------------------
+        last = end
+
+    if last < len(text):
+        html += text[last:]
+
+    return html
+
+
+# --------------------------------------
+# GRADIO UI
+# --------------------------------------
 
 with gr.Blocks() as demo:
     gr.Markdown("# Named Entity Recognition Demo using BERT")
-    gr.Markdown("### Model: **ELHACHYMI/bert-ner**")
+    gr.Markdown("### Model: ELHACHYMI/bert-ner")
 
-    # Input text
     text_input = gr.Textbox(
         label="Enter Text",
         lines=5,
-        value=example_sent,
-        placeholder="Type a sentence to analyze...",
+        value=default_example,
+        placeholder="Type a sentence to analyze."
     )
 
-    # List of examples buttons
     with gr.Row():
         for ex in example_list:
             gr.Button(ex).click(lambda x=ex: x, outputs=text_input)
 
+    mode_selector = gr.Dropdown(
+        choices=["Highlighted Mode", "Colorized HTML Mode"],
+        value="Highlighted Mode",
+        label="Output Display Mode"
+    )
+
     analyze_button = gr.Button("Run NER Analysis")
 
-    # Highlighted output
-    output = gr.HighlightedText(label="NER Result", combine_adjacent=True)
+    output_highlight = gr.HighlightedText(
+        label="NER Result (Highlighted Mode)",
+        combine_adjacent=True,
+        visible=True
+    )
+    output_html = gr.HTML(
+        label="NER Result (Colorized HTML Mode)",
+        visible=False
+    )
 
-    # Code + link output
-    code_output = gr.Code(label="How to use this model in Python", visible=True)
+    code_output = gr.Code(label="How to use this model")
     link_output = gr.Markdown()
 
-    # Initialize code + link
     demo.load(
         fn=display_model_info,
         inputs=[gr.Textbox(value=model_name, visible=False)],
-        outputs=[code_output, link_output],
+        outputs=[code_output, link_output]
     )
 
-    # Run analysis
+    def run_analysis(text, mode):
+        if mode == "Highlighted Mode":
+            return analyze_highlight(text), gr.update(visible=True), gr.update(visible=False)
+        else:
+            return gr.update(), gr.update(visible=False), colorize_text(text)
+
     analyze_button.click(
-        analyze_text, inputs=[text_input], outputs=output
+        run_analysis,
+        inputs=[text_input, mode_selector],
+        outputs=[output_highlight, output_highlight, output_html]
     )
 
 demo.launch()