import gradio as gr
import spacy
from spacy import displacy

# Load transformer-based spaCy NER pipeline
nlp = spacy.load("en_core_web_trf")  # Transformer-based model

def ner_extraction(text):
    if not text.strip():
        return "Please enter some text."
    doc = nlp(text)
    ents = [{"text": ent.text, "label": ent.label_} for ent in doc.ents]
    if not ents:
        return "No named entities found."
    return ents

# Optional: visual output
def ner_visualizer(text):
    doc = nlp(text)
    html = displacy.render(doc, style="ent", minify=True)
    return html

with gr.Blocks() as demo:
    gr.Markdown("## Named Entity Recognition using spaCy + Transformers (en_core_web_trf)")
    with gr.Tab("Extract Entities"):
        inp = gr.Textbox(label="Enter Text", lines=3, placeholder="Type a sentence...")
        out = gr.JSON(label="Named Entities")
        btn = gr.Button("Run NER")
        btn.click(ner_extraction, inputs=inp, outputs=out)

    with gr.Tab("Visualize Entities"):
        vis_inp = gr.Textbox(label="Enter Text", lines=3, placeholder="Type a sentence...")
        vis_out = gr.HTML(label="Visualization")
        vis_btn = gr.Button("Visualize")
        vis_btn.click(ner_visualizer, inputs=vis_inp, outputs=vis_out)

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