import gradio as gr
from transformers import pipeline

pipe = pipeline("text-classification", model="peter2000/xlm-roberta-base-finetuned-ecoicop")

def predict(text):
    preds = pipe(text)[0]
    return preds["label"].split('_')[1],preds["label"].split('_')[0], round(preds["score"], 5)

gradio_ui = gr.Interface(
    fn=predict,
    title="Predicting E-Coicop Product Categories from Product texts",
    description="The model is trained on product texts (shop category | text from URL | product name) from different online supermarkets in Germany, France, Austria and Italy. It predicts the corresponding ECOICOP product category (used to calculate Consumer Price Index) for food and baverages, out of 75 sub-labels on the 5-digits level of ECOICOP hierarchy.",
    inputs=[
        gr.inputs.Textbox(lines=5, label="Paste some text here"),
    ],
    outputs=[
        gr.outputs.Textbox(label="ECOICOP Label"),
        gr.outputs.Textbox(label="ECOICOP Index"),
        gr.outputs.Textbox(label="Certainty")
    ],
    examples=[
        ["Tiefkühl Eiscreme & Eiswürfel Bechereis | rewe beste wahl peanut butter eiscreme | REWE Beste Wahl Peanut Butter Eiscreme 500ml"], 
        ["epicerie-sucree | cereales chocolat fraise nat | Céréales chocolat & fraise NAT"],
        ["Pelati e passate | unknown | Mutti Polpa di Pomodoro 3 x 400 g"]
    ],
)

gradio_ui.launch(debug=True)