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

css = """

"""

ori_model = None
edit_model = None

with gr.Blocks(css=css,theme=gr.themes.Soft(text_size="sm")) as demo:
    gr.HTML("""<div style="text-align: center; margin: 0 auto;"><p><h1> Knowledge Editing</h1></div>""")

    with gr.Row():
        gr.Markdown("<p align='center'><a href='https://github.com/zjunlp/EasyEdit'>🔧https://github.com/zjunlp/EasyEdit</a></p>") 
    
    with gr.Row():
        gr.Markdown("### Large Language Models (LLMs) often suffer from knowledge cutoff or fallacy issues, meaning they are unable to recognize unseen events or generate text with incorrect facts due to outdated/noisy data. To address this, many knowledge editing methods for LLMs aim to subtly inject/edit updated knowledge or adjust undesirable behaviors, while minimizing the impact on unrelated inputs.")    
    
    
    with gr.Row():     
        prompt = gr.Textbox(label="Edit Prompt")     
        target_new = gr.Textbox(label="Edit Target New")      
    with gr.Row():
        button4clear = gr.Button("Clear")
        button4edit = gr.Button("Edit",variant="primary") 
    with gr.Row():
        examples = gr.Examples(
            examples=[
                ["Who is the architect for Toodyay Fire Station?","Wong Tung & Sons"],
                ["Who is Claire Clairmont\'s sister?","Clairmont-Mayer"],
                ["Which fictional universe is Chlorophyll Kid part of?","Image Universe"]
            ],
            examples_per_page=3,
            inputs=[prompt,target_new],
        )     
    with gr.Row():
        input_text = gr.Textbox(label="Status Information",value="Model editing may take about a minute, please be patient.")
        
    with gr.Row():
        input = gr.Textbox(label="Input Text",lines=3)
    with gr.Row():
        with gr.Column():
            button4gen_ori=gr.Textbox(label="origin output")
        with gr.Column():
            button4gen_edit=gr.Textbox(label="edited output")
    with gr.Row():
        button4gen = gr.Button("Generate",variant="primary")
   
    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()