from transformers import pipeline
import gradio as gr



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


# function for the gradio app
def fn_ner(prompt):
    return ner(prompt)


# gradio app
# 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: 1800 Antipolo Rizal Furnitown",
]

iface = gr.Interface(
    fn=fn_ner,
    inputs='text',
    outputs=gr.Label(label="Name Entity Recognition"),
    examples=[
        [ex] for ex in examples
    ],
    title='English to Tagalog translator',
    description='This demo performs language translation from English to Tagalog.',
    article='All done by Kaiku',
)

iface.launch()