import gradio as gr

def inference(model_list, cot):
    if len(model_list) != 3:
        raise gr.Error("Please choose just '3' models! Neither more nor less!")

    if cot:
        b = "CoT was also used."
    else:
        b = ""

    a = f"Hello, {model_list[0]}, {model_list[1]}, and {model_list[2]}!! {b}"
    return {
        output_msg: gr.update(visible=True),
        output_col: gr.update(visible=True),
        model1_output1: a
    }

TITLE = """<h1 align="center">LLM Agora 🗣️🏦</h1>"""

INTRODUCTION_TEXT = """
The **LLM Agora** 🗣️🏦 aims to improve the quality of open-source LMs' responses through debate & revision introduced in [Improving Factuality and Reasoning in Language Models through Multiagent Debate](https://arxiv.org/abs/2305.14325).

Do you know that? 🤔 **LLMs can also improve their responses by debating with other LLMs**! 😮 We applied this concept to several open-source LMs to verify that the open-source model, not the proprietary one, can sufficiently improve the response through discussion. 🤗
For more details, please refer to the GitHub Repository below.

You can use LLM Agora with your own questions if the response of open-source LM is not satisfactory and you want to improve the quality!
The Math, GSM8K, and MMLU Tabs show the results of the experiment, and for inference, please use the 'Inference' tab.
Please check the more specific information in [GitHub Repository](https://github.com/gauss5930/LLM-Agora)!
"""

RESPONSE_TEXT = """<h1 align="center">🤗 Here are the responses to each model!! 🤗</h1>"""

with gr.Blocks() as demo:
    gr.HTML(TITLE)
    gr.Markdown(INTRODUCTION_TEXT)
    with gr.Column():
        with gr.Tab("Inference"):
            with gr.Row():
                with gr.Column():
                    model_list = gr.CheckboxGroup(["Llama2", "Alpaca", "Vicuna", "Koala", "Falcon", "Baize", "WizardLM", "Orca", "phi-1.5"], label="Model Selection", info="Choose 3 LMs to participate in LLM Agora.", type="value")
                    cot = gr.Checkbox(label="CoT", info="Do you want to use CoT for inference?")
                with gr.Column():
                    API_KEY = gr.Textbox(label="OpenAI API Key", value="", info="Please fill in your OpenAI API token.", placeholder="sk..", type="password")
                    auth_token = gr.Textbox(label="Huggingface Authentication Token", value="", info="Please fill in your HuggingFace Authentication token.", placeholder="hf..", type="password")
            with gr.Column():
                question = gr.Textbox(value="", info="Please type your question!", placeholder="")
                output = gr.Textbox()
            submit = gr.Button("Submit")
            with gr.Row(visible=False) as output_msg:
                gr.HTML(RESPONSE_TEXT)

            with gr.Row(visible=False) as output_col:
                with gr.Column():
                    model1_output1 = gr.Textbox(label="1️⃣ model's initial response")
                    model1_output2 = gr.Textbox(label="1️⃣ model's revised response")
                    model1_output3 = gr.Textbox(label="1️⃣ model's final response")
                with gr.Column():
                    model2_output1 = gr.Textbox(label="2️⃣ model's initial response")
                    model2_output2 = gr.Textbox(label="2️⃣ model's revised response")
                    model2_output3 = gr.Textbox(label="2️⃣ model's final response")
                with gr.Column():
                    model2_output1 = gr.Textbox(label="3️⃣ model's initial response")
                    model2_output2 = gr.Textbox(label="3️⃣ model's revised response")
                    model2_output3 = gr.Textbox(label="3️⃣ model's final response")
            

        with gr.Tab("Math"):
            output_math = gr.Textbox()
        with gr.Tab("GSM8K"):
            output_gsm = gr.Textbox()
        with gr.Tab("MMLU"):
            output_mmlu = gr.Textbox()

    submit.click(inference, [model_list], [output_msg, output_col, model1_output1])

demo.launch(debug=True)