from transformers import pipeline
import gradio as gr

def coding(model, text, codetext):
    classifier = pipeline("zero-shot-classification", model=model)
    codelist = codetext.split(';')
    output = classifier(text, codelist, multi_label=True)
    return output

iface = gr.Interface(
    fn=coding, 
    inputs=[
        gr.Radio(
            [
                "facebook/bart-large-mnli",
                "MoritzLaurer/multilingual-MiniLMv2-L6-mnli-xnli",
                "MoritzLaurer/mDeBERTa-v3-base-xnli-multilingual-nli-2mil7",
                "MoritzLaurer/mDeBERTa-v3-base-mnli-xnli",
                "MoritzLaurer/deberta-v3-large-zeroshot-v2.0",
                #"joeddav/xlm-roberta-large-xnli"
            ],
            #min_width=200,
            #scale=2,
            value="facebook/bart-large-mnli",
            label="Model"
        ),
        gr.TextArea(
            label='Comment',
            value='感覺性格溫和，適合香港人，特別係亞洲人的肌膚，不足之處就是感覺很少有優惠，價錢都比較貴'
        ),
        gr.Textbox(
            label='Code list (colon-separated)',
            value='非常好/很好/好滿意;價錢合理/實惠/不太貴/親民/價格適中/價格便宜/價錢大眾化;價錢貴/不合理/比日本台灣貴/可以再平d'
        )
    ],
    outputs=[
        #gr.Textbox(label='Result')
        gr.JSON()
        #gr.BarPlot()
    ],
    title="NuanceTree Coding Test",
    description="Test Zero-Shot Classification",
    allow_flagging='never'
)

iface.launch()