from transformers import pipeline
import gradio as gr
import pandas as pd

# Load NER model
ner_model = pipeline(
    "ner",
    model="dslim/bert-base-NER",
    aggregation_strategy="simple"
)

def create_df():
    return pd.DataFrame(
        columns=["Entity", "Confidence (%)"]
    )

def ner_inference(text):
    entities = ner_model(text)

    per_rows, org_rows, loc_rows = [], [], []

    for ent in entities:
        row = [ent["word"], round(ent["score"] * 100, 2)]

        if ent["entity_group"] == "PER":
            per_rows.append(row)
        elif ent["entity_group"] == "ORG":
            org_rows.append(row)
        elif ent["entity_group"] == "LOC":
            loc_rows.append(row)

    df_per = pd.DataFrame(per_rows, columns=["Person", "Confidence (%)"]) if per_rows else create_df()
    df_org = pd.DataFrame(org_rows, columns=["Organization", "Confidence (%)"]) if org_rows else create_df()
    df_loc = pd.DataFrame(loc_rows, columns=["Location", "Confidence (%)"]) if loc_rows else create_df()

    return df_per, df_org, df_loc

# Gradio UI
interface = gr.Interface(
    fn=ner_inference,
    inputs=gr.Textbox(
        lines=5,
        placeholder="Enter text here...",
        label="Input Text"
    ),
    outputs=[
        gr.Dataframe(label="👤 Persons", interactive=False),
        gr.Dataframe(label="🏢 Organizations", interactive=False),
        gr.Dataframe(label="📍 Locations", interactive=False),
    ],
    title="Named Entity Recognition (NER)",
    description="NER results grouped by entity type for better readability and usability.",
    theme="dark"
)

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