import gradio as gr
from utils import *
from transformers import pipeline

css = """
button {
    background-color: #673AB7; /* 设置按钮背景为紫色 */
    color: white; /* 设置按钮文字颜色为白色 */
    border-radius: 8px; /* 设置按钮边角为圆角 */
    padding: 10px 20px; /* 设置按钮内边距 */
    border: none; /* 移除按钮边框 */
}
"""

ori_model = None
edit_model = None

with gr.Blocks(css=css) as demo:
    gr.Markdown("# Model Output Editor")
    gr.Markdown("This interface takes your input, shows the output of the source model, and then the edited model's output.")
    with gr.Row():
        with gr.Column():
            with gr.Row():     
                prompt = gr.Textbox(label="Input Prompt",lines=4)
            with gr.Row():
                target_new = gr.Textbox(label="Input Target New")
            with gr.Row():
                button4clear = gr.Button("Clear")
                button4edit = gr.Button("Edit")
            with gr.Row():
                input_text = gr.Label(label="Status Information",value="The editing process may take up to 30 seconds. Please be patient.")
        with gr.Column():
            with gr.Row():
                input = gr.Textbox(label="Input Text")
            with gr.Row():
                button4gen = gr.Button("Generate")
            with gr.Row():
                button4gen_ori=gr.Label(label="origin output")
            with gr.Row():
                button4gen_edit=gr.Label(label="edited output")
                
    button4clear.click(lambda: ("", ""), outputs=[prompt,target_new])
    button4edit.click(fn=edit, inputs=[prompt,target_new], outputs=input_text)
    button4gen.click(fn=generate, inputs=input, outputs=[button4gen_ori,button4gen_edit])


demo.launch()