from transformers import pipeline
import gradio as gr

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

def extract_named_entities(text):
    # Use the NER model to extract named entities
    entities = ner_model(text)
    # Format the output for better readability
    return [{"Entity": ent["entity_group"], "Text": ent["word"], "Score": round(ent["score"], 3)} for ent in entities]

# Define the Gradio interface
iface = gr.Interface(
    fn=extract_named_entities,
    inputs=gr.Textbox(lines=5, label="Input Text"),
    outputs=gr.Dataframe(headers=["Entity", "Text", "Score"], label="Named Entities"),
    title="Named Entity Recognition",
    description="Input some text and get the named entities (like names, locations, organizations).",
)

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