import gradio as gr
from transformers import pipeline

# Load NER pipeline
ner_pipeline = pipeline(
    "ner",
    model="dbmdz/bert-large-cased-finetuned-conll03-english",
    aggregation_strategy="simple"
)

def ner_analysis(text):
    try:
        results = ner_pipeline(text)

        if not results:
            return "No entities found."

        output = []
        for r in results:
            output.append(
                f"{r['word']} → {r['entity_group']} (Score: {round(r['score'], 4)})"
            )

        return "\n".join(output)

    except Exception as e:
        return f"Error: {str(e)}"

# UI
with gr.Blocks() as app:
    gr.Markdown("# 🧠 Named Entity Recognition (NER)")
    gr.Markdown("Detect persons, organizations, and locations from text.")

    text_input = gr.Textbox(
        label="Enter text",
        value="My name is Sylvain and I work at Hugging Face in Brooklyn."
    )

    output = gr.Textbox(label="Entities")

    submit_btn = gr.Button("Analyze")

    submit_btn.click(
        fn=ner_analysis,
        inputs=text_input,
        outputs=output
    )

# Launch
if __name__ == "__main__":
    app.launch()