from transformers import pipeline
import gradio as gr
import pandas as pd


def coding(model, text, codetext):
    classifier = pipeline("zero-shot-classification", model=model)
    codelist = codetext.split(';')
    output = classifier(text, codelist, multi_label=True)
    # keys = output.labels
    # values = output.scores
    keys = output['labels']
    values = output['scores']
    my_dict = {k: v for k, v in zip(keys, values)}
    return [my_dict, output]
    
def upload_code_list(file):
    df = pd.read_excel(file.name, sheet_name='code')
    joined_data = ';'.join(df['label'].astype(str))
    return joined_data


css = """
h2.svelte-1pq4gst{display:none}
"""

demo = gr.Blocks(css=css)

with demo:
    gr.Markdown(
        """
        # NuanceTree
        # Coding Test Program
        """
    )
     
    with gr.Row():
        with gr.Column():
            select_model = 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",
                ],
                value="MoritzLaurer/multilingual-MiniLMv2-L6-mnli-xnli",
                label="Model"
            )
            comment_text = gr.TextArea(
                label='Comment',
                value='感覺性格溫和，特別係亞洲人的肌膚，不足之處就是感覺很少有優惠，價錢都比較貴'
            )
            codelist_text = gr.Textbox(
                label='Code list (colon-separated)',
                value='非常好;很好;好滿意;價錢合理;實惠'
            )
            with gr.Row():
                clear_codelist_btn = gr.ClearButton(value="Clear Code List")
                clear_codelist_btn.click(lambda: None, outputs=[codelist_text]) 
                upload_btn = gr.UploadButton(
                    label="Upload",
                    variant='primary'
                )
                upload_btn.upload(upload_code_list, upload_btn, codelist_text)
            run_btn = gr.Button(
                value="Submit",
                variant='primary'
            )
        with gr.Column():
            result_label = gr.Label(show_label=False)
            result_text = gr.JSON()
            
    run_btn.click(coding, [select_model, comment_text, codelist_text], [result_label, result_text], scroll_to_output=True)

if __name__ == "__main__":
    demo.launch()