import gradio as gr
from transformers import pipeline

# Initialize Hugging Face NER pipeline
ner_pipeline = pipeline("ner", model="HooshvareLab/bert-fa-base-uncased", aggregation_strategy="simple")

def split_name(full_name):
    entities = ner_pipeline(full_name)
    prefix = ""
    first_name = ""
    last_name = ""
    for entity in entities:
        label = entity['entity_group']
        word = entity['word']
        if label.lower() == 'prefix':
            prefix += word + " "
        elif label.lower() == 'firstname':
            first_name += word + " "
        elif label.lower() == 'lastname':
            last_name += word + " "
    return prefix.strip(), first_name.strip(), last_name.strip()

# Gradio interface
demo = gr.Interface(
    fn=split_name,
    inputs=gr.Textbox(label="Full Name (Persian)", placeholder="Enter your full name"),
    outputs=[
        gr.Textbox(label="Prefix"),
        gr.Textbox(label="First Name"),
        gr.Textbox(label="Last Name"),
    ],
)

demo.launch()