import gradio as gr
from transformers import pipeline

ner = pipeline(task="token-classification", model="IsmaelMousa/modernbert-ner-conll2003", aggregation_strategy="max")

def extract(text):
    """
    Extract named entities from text
    
    :param text: input text
    :return: formatted output and highlighted text
    """
    if not text.strip(): return "Please enter some text to analyze.", []
    
    results = ner(text)
    output  = ""
    
    for entity in results:
        word    = entity["word"]
        label   = entity["entity_group"]
        score   = entity["score"]
        output += f"**{word}** → {label} (confidence: {score:.2%})\n"
    
    if not output: output = "No named entities found in the text."
    
    highlighted = []
    last        = 0
    
    for entity in results:
        start = entity["start"]
        end   = entity["end"]
        
        if start > last: highlighted.append((text[last:start], None))
        
        highlighted.append((text[start:end], entity["entity_group"]))
        last = end
    
    if last < len(text): highlighted.append((text[last:], None))
    
    return output, highlighted if highlighted else [(text, None)]

examples = [["Hi, I'm Ismael Mousa from Palestine working for NVIDIA inc."]             ,
            ["The conference was held in Paris by the World Health Organization."]      ,
            ["John Smith joined Microsoft in Seattle office."]                          ,
            ["IBM announced new investments in India last year. Wrote by Gholam Ghanni"],]

with gr.Blocks(title="Named Entity Recognition") as demo:
    gr.Markdown(
        """
        # 🏷️ Named Entity Recognition
        Extract named entities (persons, organizations, locations) from text using ModernBERT.
        
        **Model:** [IsmaelMousa/modernbert-ner-conll2003](https://huggingface.co/IsmaelMousa/modernbert-ner-conll2003)
        """
    )
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="Input Text", placeholder="Enter text to analyze...", lines=5)
            submit_btn = gr.Button("Extract Entities", variant="primary")
        
        with gr.Column():
            output_text      = gr.Markdown(label="Detected Entities")
            highlighted_text = gr.HighlightedText(label="Highlighted Text", combine_adjacent=True, show_legend=True)
    
    gr.Examples(examples=examples, inputs=input_text, outputs=[output_text, highlighted_text], fn=extract, cache_examples=False)
    
    submit_btn.click(fn=extract, inputs=input_text, outputs=[output_text, highlighted_text])
    input_text.submit(fn=extract, inputs=input_text, outputs=[output_text, highlighted_text])

if __name__ == "__main__": demo.launch()