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

css = """

"""

ori_model = None
edit_model = None
# input=None

def slowly_reverse(word, progress=gr.Progress()):
    progress(0, desc="Starting")
    time.sleep(1)
    progress(0.05)
    new_string = ""
    for letter in progress.tqdm(word, desc="Editing"):
        time.sleep(0.25)
        new_string = letter + new_string
    return new_string

with gr.Blocks(css=css,theme=gr.themes.Soft(text_size="sm")) as demo:
    with gr.Row(equal_height=True):
        gr.HTML(
                """
                <div style="display: flex; flex-direction: column; align-items: center;">
                    <h1>🔧EasyEdit: An Easy-to-use Knowledge Editing Framework for Large Language Models</h1>
                    
                    <p>
                    📑[<a href="https://huggingface.co/papers/2308.07269">Paper</a>]
                    👨‍💻[<a href="https://github.com/zjunlp/EasyEdit" target="_blank"><span class="icon"><i class="fab fa-github"></i></span>Code</a>]
                    📄[<a href="https://zjunlp.gitbook.io/easyedit">Docs</a>]
                    🤗[<a href="https://huggingface.co/spaces/zjunlp/EasyEdit" target="_blank">Demo</a>]
                    [<a href="https://arxiv.org/abs/2211.11031">via GRACE</a>]
                    </p>
                </div>
                """
)
  # 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("#### Knowledge editing aims to subtly inject/edit updated knowledge or adjust undesirable behaviors, while minimizing the impact on unrelated inputs.")    
    with gr.Accordion("Expiation", open=False):
        gr.Markdown(
            """
            Edit Steps: the number of times a layer is trained in the GRACE method.
            """
        )
        gr.Markdown(
            """
            Replacement: the optimization strategy during fine-tuning.
            """
        )
        gr.Markdown(
            """
            Reliability Evaluation: the optimization strategy during fine-tuning.
            """
        )       
        gr.Markdown(
            """
            Reliability Evaluation: the assessment of whether the target edit can be accomplished.
            """
        )          
        gr.Markdown(
            """
            Locality Evaluation: the assessment of whether unrelated content has been affected..
            """
        )   
    
    with gr.Row():     
        prompt = gr.Textbox(label="Edit Prompt")     
        target_new = gr.Textbox(label="Edit Target New")
    with gr.Row():
        num_steps = gr.Slider(10, 100, value=40, step=1, label='Edit Steps')
        replacement = gr.Dropdown(
            choices=["replace_last", "replace_all", "replace_prompt"],
            value="replace_last",
            label="Replacement",
        )
    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():
        gr.HTML(
            """
            <h3>Reliability Evaluation</h3>
            """
        )
    with gr.Row():
        input = gr.Textbox(label="Input Text")
    with gr.Row():
        with gr.Column():
            button4gen_ori=gr.HighlightedText(
                label="original output",
                combine_adjacent=True,
                show_legend=False,
                color_map={"output": "yellow"},
            )
        with gr.Column():
            button4gen_edit=gr.HighlightedText(
                label="edited output",
                combine_adjacent=True,
                show_legend=False,
                color_map={"output": "yellow"},
            )
    with gr.Row():
        button4gen = gr.Button("Generate",variant="primary")

    with gr.Row():
        gr.HTML(
            """
            <h3>Locality Evaluation</h3>
            """
        )
    with gr.Row():
        loc_input = gr.Dropdown(
            choices=[
                "who sang the theme song for laverne and shirley",
                "when does the last episode of adventure time air",
                "who plays alec ramsay in the black stallion",
                "where did an independence movement occur because of the congress of vienna",
                "where is the ucla usc game being played"
                ],
            value="where is the ucla usc game being played",
            label="Unrelated Input Text",
        )
    with gr.Row():
        with gr.Column():
            button4gen_loc_ori=gr.HighlightedText(
                label="original output",
                combine_adjacent=True,
                show_legend=False,
                color_map={"output": "green"},
            )
        with gr.Column():
            button4gen_loc_edit=gr.HighlightedText(
                label="edited output",
                combine_adjacent=True,
                show_legend=False,
                color_map={"output": "green"},
            )
    with gr.Row():
        button4locgen = gr.Button("Generate",variant="primary")
   
    button4clear.click(lambda: ("", ""), outputs=[prompt,target_new])
    button4edit.click(fn=edit, inputs=[prompt,target_new, num_steps, replacement], outputs=input)
    button4gen.click(fn=generate, inputs=[input, target_new], outputs=[button4gen_ori, button4gen_edit])
    button4locgen.click(fn=generate, inputs=loc_input, outputs=[button4gen_loc_ori, button4gen_loc_edit])


    with gr.Accordion("Citation", open=False):
        gr.Markdown(
            """
            ```bibtex
            @misc{wang2023easyedit,
                title={EasyEdit: An Easy-to-use Knowledge Editing Framework for Large Language Models}, 
                author={Peng Wang and Ningyu Zhang and Xin Xie and Yunzhi Yao and Bozhong Tian and Mengru Wang and Zekun Xi and Siyuan Cheng and Kangwei Liu and Guozhou Zheng and Huajun Chen},
                year={2023},
                eprint={2308.07269},
                archivePrefix={arXiv},
                primaryClass={cs.CL}
            }
            ```
            """
        )


demo.launch()