from transformers import pipeline
import gradio as gr

# Load the NER pipeline
ner = pipeline(
    "ner",
    model='kaiku03/bert-base-NER-finetuned_custom_complain_dataset_NER9',
    aggregation_strategy="simple"
)

# Function for the Gradio app
def fn_ner(prompt):
    entities = ner(prompt)
    entity_info = "\n".join([f"{entity['entity_group']}: {entity['word']}" for entity in entities])
    return entity_info

# Define example inputs and outputs
examples = [
    "Subject: Defective Date: 08-13-2023 Product: XXX speaker Location: 456 Sound Avenue, Audiotown",
    "Subject: Dirty Date: 08-10-2023 Product: UVW Television Location: 567 Willow Lane, Mediatown",
    "Subject: Missing Parts Date: 08-10-2023 Product: XXX Furniture Set Location: Antipolo Rizal",
]

# Create Gradio interface
iface = gr.Interface(
    fn=fn_ner,
    inputs='text',
    outputs='text',
    examples=[[ex] for ex in examples],
    title='Named Entity Recognition',
    description='This demo performs named entity recognition (NER) on our custom made dataset. This dataset consist of very small training and testing samples resulting a very limited data our model to learn. ',
    article='All done by Kaiku'
)

# Launch the interface
iface.launch()