import test_
import mdtex2html
import gradio as gr
from transformers import AutoModel, AutoTokenizer, AutoConfig

def postprocess(self, y):
    if y is None:
        return []
    for i, (message, response) in enumerate(y):
        y[i] = (
            None if message is None else mdtex2html.convert((message)),
            None if response is None else mdtex2html.convert(response),
        )
    return y

gr.Chatbot.postprocess = postprocess

def parse_codeblock(text):
    lines = text.split("\n")
    for i, line in enumerate(lines):
        if "```" in line:
            if line != "```":
                lines[i] = f'<pre><code class="{lines[i][3:]}">'
            else:
                lines[i] = '</code></pre>'
        else:
            if i > 0:
                lines[i] = "<br/>" + line.replace("<", "&lt;").replace(">", "&gt;")
    return "".join(lines)


def predict(input, chatbot, history):
        
    # map_ = {'1': '2', '2': '2', '3': '4'}

    chatbot.append((input, ""))
    response = test_.map_(input_text=input)
    # print(response)
    # print(response)
    chatbot[-1] = (parse_codeblock(input), parse_codeblock(response))
    return chatbot, history

def reset_user_input():
    return gr.update(value='')

def reset_state():
    return [], []

with gr.Blocks() as demo:
    gr.HTML("""<h1 align="center">错误分类</h1>""")
    # gr.Markdown("Start typing below and then click **Run** to see the output.")

    chatbot = gr.Chatbot()
    with gr.Row():
        with gr.Column(scale=4):
            with gr.Column(scale=12):
                # user_input = gr.Textbox(show_label=False, placeholder="Input...", lines=10).style(
                #     container=False)
                user_input = gr.Textbox(show_label=False, placeholder="Input...", lines=10, container=False)
            with gr.Column(min_width=32, scale=1):
                submitBtn = gr.Button("Submit", variant="primary")
        with gr.Column(scale=1):
            emptyBtn = gr.Button("Clear History")
            # max_length = gr.Slider(0, 4096, value=2048, step=1.0, label="Maximum length", interactive=True)
            # top_p = gr.Slider(0, 1, value=0.7, step=0.01, label="Top P", interactive=True)
            # temperature = gr.Slider(0, 1, value=0.95, step=0.01, label="Temperature", interactive=True)

    history = gr.State([])

    submitBtn.click(predict, [user_input, chatbot, history], [chatbot, history],
                    show_progress=True)
    submitBtn.click(reset_user_input, [], [user_input])

    emptyBtn.click(reset_state, outputs=[chatbot, history], show_progress=True)

demo.queue().launch(share=False, inbrowser=True)