import gradio as gr
from transformers import pipeline

def token_classification(sentence):
    model_checkpoint = "AirrStorm/bert-ner-finetuned"
    token_classifier = pipeline(
        "token-classification", model=model_checkpoint, aggregation_strategy="simple"
    )

    # Get predictions
    results = token_classifier(sentence)

    # Clean and format the results
    formatted_results = "Named Entities Found:\n"
    for result in results:
        entity = result['entity_group']
        word = result['word']
        formatted_results += f"- {entity}: {word}\n"
    
    return formatted_results

examples = [
        ["Elon Musk founded SpaceX in California."],
        ["Cristiano Ronaldo plays for Al Nassr in Saudi Arabia."],
        ["Apple Inc. was founded by Steve Jobs. It is headquartered in Cupertino, California."],
        ["Albert Einstein developed the theory of relativity in Germany."],
        ["Harry Potter studied at Hogwarts in the United Kingdom."]
                ]

# Gradio interface
demo = gr.Interface(
    fn=token_classification, 
    inputs=gr.Textbox(label="Enter your sentence here", placeholder= "Enter a sentence to extract named entities (e.g., 'Lionel Messi plays for Inter Miami in the United States.')", lines=3),
    outputs=gr.Textbox(label="Extracted Named Entities", placeholder="The extracted entities will appear here", lines=6),
    title="Named Entity Recognition",
    description="Extract and classify named entities like people, organizations, and locations from the input text.",
    theme="dark-grass",  # Optional: Use a clean and minimal theme
    allow_flagging="never",  # Optional: Disable flagging to reduce distractions
    examples=examples,  # Optional: Add examples to make it easier for users
)
demo.launch()