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)
    return output

def upload_code_list(file):
    df = pd.read_excel(file.name, sheet_name='code')
    # Join the data in column B using ";" as the delimiter
    joined_data = ';'.join(df['label'].astype(str))
    #file_paths = [file.name for file in files]
    
    return joined_data

demo = gr.Blocks()

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",
                    #"joeddav/xlm-roberta-large-xnli"
                ],
                #min_width=200,
                #scale=2,
                value="facebook/bart-large-mnli",
                label="Model"
            )
            comment_text = gr.TextArea(
                label='Comment',
                value='感覺性格溫和，適合香港人，特別係亞洲人的肌膚，不足之處就是感覺很少有優惠，價錢都比較貴'
            )
            # upload_btn = gr.UploadButton(
            #     label="Upload the code list file"
            # )
            codelist_text = gr.Textbox(
                label='Code list (colon-separated)',
                value='非常好/很好/好滿意;價錢合理/實惠/不太貴/親民/價格適中/價格便宜/價錢大眾化;價錢貴/不合理/比日本台灣貴/可以再平d'
            )
            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 the code list file",
                    variant='primary'
                )
                upload_btn.upload(upload_code_list, upload_btn, codelist_text)
            run_btn = gr.Button(
                value="Submit",
                variant='primary'
            )
        with gr.Column():
            result_text = gr.JSON()
    run_btn.click(coding, [select_model, comment_text, codelist_text], result_text, scroll_to_output=True)

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