app.py

import gradio as gr
from transformers import pipeline

# Load model
model_name = "Salommee/bert-finetuned-ner"
ner_pipeline = pipeline(
    "token-classification",
    model=model_name,
    aggregation_strategy="simple"
)

def analyze_text(text):
    """Extract named entities and format nicely"""
    if not text.strip():
        return "No text provided", []
    
    results = ner_pipeline(text)
    
    if not results:
        return text, []
    
    # Format as list of tuples: (text, label) for each entity
    # This works with HighlightedText in Gradio
    highlighted = []
    last_end = 0
    
    for entity in results:
        # Add text before entity
        if entity['start'] > last_end:
            highlighted.append((text[last_end:entity['start']], None))
        
        # Add entity with label
        highlighted.append((entity['word'], entity['entity_group']))
        last_end = entity['end']
    
    # Add remaining text
    if last_end < len(text):
        highlighted.append((text[last_end:], None))
    
    return highlighted

# Create interface
with gr.Blocks(title="Named Entity Recognition") as demo:
    gr.Markdown("# 🏷️ Named Entity Recognition")
    gr.Markdown("Extract persons, organizations, locations, and miscellaneous entities from text")
    
    with gr.Row():
        input_text = gr.Textbox(
            lines=5,
            placeholder="Enter text to analyze...",
            label="Input Text"
        )
    
    with gr.Row():
        submit_btn = gr.Button("🔍 Analyze", variant="primary")
    
    with gr.Row():
        output = gr.HighlightedText(
            label="Detected Entities",
            color_map={
                "PER": "lightblue",
                "ORG": "lightgreen",
                "LOC": "yellow",
                "MISC": "pink"
            }
        )
    
    gr.Examples(
        examples=[
            ["Apple CEO Tim Cook announced new products in California."],
            ["The United Nations headquarters is located in New York City."],
            ["Microsoft acquired GitHub for $7.5 billion in 2018."],
            ["Barack Obama was born in Hawaii and became the 44th President."]
        ],
        inputs=input_text
    )
    
    submit_btn.click(fn=analyze_text, inputs=input_text, outputs=output)

demo.launch(share=True)