import gradio as gr
from simpletransformers.ner import NERModel 
import string

labels = ["O", "B-FOOD_QUANTITY", "B-FOOD_SIZE", "B-FOOD", "I-FOOD", "B-FOOD_INGREDIENTS", "I-FOOD_INGREDIENTS", "B-DRINK_SIZE", "B-DRINK_QUANTITY", "B-DRINK", "B-PAYMENT", "I-PAYMENT", "B-DELIVERY_ADDRESS", "I-DRINK_SIZE", "I-DRINK", "I-FOOD_SIZE", "I-DELIVERY_ADDRESS"]

model = NERModel(
        "roberta", 
        "bgk/berteng", labels=labels,
        use_cuda=False,
        ignore_mismatched_sizes=True
    )

examples=[['I want two hamburgers and one sprite and one milkshake, send it to my workplace.' ], [' I want to order two large pizzas, two medium coke, send it to my home, I will pay with cash.' ]]

def ner(text):
    trans_table = text.maketrans('', '', string.punctuation)
    text = text.translate(trans_table)
    text=text.lower()
    
    prediction, model_output = model.predict([text])
    
    filtered_output = (({v: k} for d in sublist for k, v in d.items() if (v.startswith("B-") or v.startswith("I-"))) for sublist in prediction)
    entities = []
    for sublist in filtered_output:
        for d in sublist:
            for k, v in d.items():
                label = k.split("-")[1]
                entities.extend([(label, v)])
              
    return  entities # prediction
    
demo = gr.Interface(ner,
             gr.Textbox(placeholder="Enter your sentences here..."),   
             gr.HighlightedText(),
             examples=examples)


demo.launch()