from transformers import pipeline
import gradio as gr

# Load the NER pipeline
try:
    ner_model = pipeline("ner", model="dslim/bert-base-NER", aggregation_strategy="simple")
    print("Model loaded successfully.")
except Exception as e:
    ner_model = None
    print(f"Error loading model: {e}")

def extract_named_entities(text):
    if ner_model is None:
        return [["Error", "Model not loaded", 0.0]]

    if not text.strip():
        return [["Error", "No input provided", 0.0]]

    try:
        entities = ner_model(text)
        # Convert list of dictionaries to list of lists for Gradio compatibility
        return [[ent["entity_group"], ent["word"], round(ent["score"], 3)] for ent in entities]
    except Exception as e:
        return [["Error", str(e), 0.0]]

# 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).",
)

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