app.py

import gradio as gr
from utils import *
from transformers import pipeline
import random
import torch
import numpy as np
seed=0
random.seed(seed)
torch.manual_seed(seed)
np.random.seed(seed)
torch.cuda.manual_seed_all(seed)

ori_model = None
edit_model = None

css = '''
'''
# 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

def single_edit_tab():
    with gr.Row():     
        prompt = gr.Textbox(label="Edit Prompt")     
        target_new = gr.Textbox(label="Edit Target New")
    with gr.Row():
        edit_alg = gr.Dropdown(
            choices=['ROME', 'WISE', 'GRACE'],
            value='WISE',
            label="Edit Algorithm",
        )
        num_steps = gr.Slider(10, 100, value=40, step=1, label='Edit Steps')
        edit_lr = gr.Dropdown(
            choices=[0.1, 0.5, 1.0],
            value=1.0,
            label="Edit LR (learning rate)",
        )
    with gr.Row():
        examples = gr.Examples(
            examples=[
                ["Who is the architect for Toodyay Fire Station?","Wong Tung & Sons"],
                ["What company makes Springfield Armory XDM?","Messerschmitt"],
                ["Which fictional universe is Chlorophyll Kid part of?","Image Universe"]
            ],
            examples_per_page=3,
            inputs=[prompt,target_new],
        )     
    with gr.Row():
        button4clear = gr.Button("Clear")
        button4edit = gr.Button("Edit",variant="primary") 
    # 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>Evaluation</h3>
            """
        )
    with gr.Row():
        gr.HTML(
            """
            <h4>Reliability</h4>
            """
        )
    # with gr.Row():
    #     input = gr.Textbox(label="Input Text")
    #     target = gr.Textbox(label="Input Answer", visible=False)
    target = gr.Textbox(label="Answer", visible=False)
    with gr.Row():
        input = gr.Textbox(label="Edit Prompt")
        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():
        gr.HTML(
            """
            <h4>Generalization</h4>
            """
        )
    with gr.Row():
        para_input = gr.Textbox(label="Paraphrase Prompt")
        with gr.Column():
            button4gen_para_ori=gr.HighlightedText(
                label="original output",
                combine_adjacent=True,
                show_legend=False,
                color_map={"output": "blue"},
            )
        with gr.Column():
            button4gen_para_edit=gr.HighlightedText(
                label="edited output",
                combine_adjacent=True,
                show_legend=False,
                color_map={"output": "blue"},
            )
    with gr.Row():
        examples = gr.Examples(
            examples=[
                ["Who is the architect for Toodyay Fire Station?", "Who was responsible for the planning of the Toodyay Fire Station", "Wong Tung & Sons"],
                ["What company makes Springfield Armory XDM?", "Which company produced Springfield Armory XDM?", "Messerschmitt"],
                ["Which fictional universe is Chlorophyll Kid part of?", "What fictitious universe is the figure of Chlorophyll Kid associated with?", "Image Universe"]
            ],
            examples_per_page=3,
            inputs=[input, para_input, target],
            label='Evaluation Examples'
        )  
    with gr.Row():
        button4gen = gr.Button("Generate",variant="primary")

    with gr.Row():
        gr.HTML(
            """
            <h4>Locality</h4>
            """
        )
    with gr.Row():
        loc_input = gr.Dropdown(
            choices=[
                "nq question: where does the phrase good bye felicia come from",
                "nq question: which best describes timbuktu under the mali empire",
                "nq question: where do the question marks go in spanish",
                "nq question: who replaces the vice president in the senate",
                "nq question: active transport performs which function in a cell"
                ],
            value="nq question: which best describes timbuktu under the mali empire",
            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(fn=clear, outputs=[prompt,target_new])
    button4edit.click(fn=edit, inputs=[edit_alg, prompt,target_new, num_steps, edit_lr], outputs=[input, target])
    button4gen.click(fn=union_generate, inputs=[input, para_input, target, edit_alg], outputs=[button4gen_ori, button4gen_edit, button4gen_para_ori, button4gen_para_edit])
    # button4gen.click(fn=generate, inputs=[para_input, target, edit_alg], outputs=[button4gen_para_ori, button4gen_para_edit])
    button4locgen.click(fn=generate, inputs=loc_input, outputs=[button4gen_loc_ori, button4gen_loc_edit])

def continuous_edit_tab():
    with gr.Accordion("Overview", open=True):
        gr.Markdown(
            """
            - [Continuous Knowledge Editing](https://arxiv.org/abs/2405.14768) is defined as multiple edits on the same model. We provide all the editing examples (10 in total) in the 'Evaluation Examples' section.

            >- **Note 1**: ❗️❗️❗️ In the cold start phase, we have already continuously edited the first 6 examples, so you can proceed directly to the evaluation tests.

            >- **Note 2**: The models edited by WISE and GRACE are the **same but independent** of each other. You need to switch the `Edit Algorithm` at the top for editing/evaluation.
            """
        )
    with gr.Row():
        # edit_alg = gr.Dropdown(
        #     choices=['WISE', 'GRACE'],
        #     value='WISE',
        #     label="Edit Algorithm",
        # )
        edit_alg = gr.Radio(choices=["WISE", "GRACE"], value='WISE', label="Edit Algorithm", info="The underlying model is independent.")
    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')
        edit_lr = gr.Dropdown(
            choices=[0.1, 0.5, 1.0],
            value=1.0,
            label="Edit LR (learning rate)",
        )
    with gr.Row():
        examples = gr.Examples(
            examples=[
                ["What is the date of birth for Christoph von Stadion?", "12 April 1809"],
                ["What medical condition killed Ramesses V?", "esses IV"],
                ["What voice type is Nellie Briercliffe?", "mezzo-oprano"],
                ["What network is 1000 Ways to Die associated with?", "The CW"]
            ],
            examples_per_page=4,
            inputs=[prompt,target_new],
        )     
    with gr.Row():
        button4edit = gr.Button("Edit",variant="primary") 
    # 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>Evaluation</h3>
            """
        )
    with gr.Row():
        gr.HTML(
            """
            <h4>Reliability</h4>
            """
        )
    # with gr.Row():
    #     input = gr.Textbox(label="Input Text")
    #     target = gr.Textbox(label="Input Answer", visible=False)
    target = gr.Textbox(label="Answer", visible=False)
    with gr.Row():
        input = gr.Textbox(label="Edit Prompt")
        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():
        gr.HTML(
            """
            <h4>Generalization</h4>
            """
        )
    with gr.Row():
        para_input = gr.Textbox(label="Paraphrase Prompt")
        with gr.Column():
            button4gen_para_ori=gr.HighlightedText(
                label="original output",
                combine_adjacent=True,
                show_legend=False,
                color_map={"output": "blue"},
            )
        with gr.Column():
            button4gen_para_edit=gr.HighlightedText(
                label="edited output",
                combine_adjacent=True,
                show_legend=False,
                color_map={"output": "blue"},
            )
    with gr.Row():
        examples = gr.Examples(
            examples=[
                ["Who is the architect for Toodyay Fire Station?", "Who was responsible for the planning of the Toodyay Fire Station", "Wong Tung & Sons"],
                ["What company makes Springfield Armory XDM?", "Which company produced Springfield Armory XDM?", "Messerschmitt"],
                ["Which fictional universe is Chlorophyll Kid part of?", "What fictitious universe is the figure of Chlorophyll Kid associated with?", "Image Universe"],
                ["What year did Sunnyside Hospital cease to exist?", "What year was the end of Sunnyside Hospital?", "1962"],
                ["Which designer was responsible for Holmenkollen Chapel?", "Which designer is responsible for Holmenkollen Chapel?", "Inigo Jones"],
                ["What piece of fiction does Jack Harkness appear in?", "What fictional work does Jack Harkness exist in?", "Lost"],
                ["What is the date of birth for Christoph von Stadion?", "What is Christoph von Stadion's birth date?", "12 April 1809"],
                ["What medical condition killed Ramesses V?", "What kind of disease killed Ramesses V?", "esses IV"],
                ["What voice type is Nellie Briercliffe?", "Which was the voice type that Nellie Briercliffe had?", "mezzo-oprano"],
                ["What network is 1000 Ways to Die associated with?", "The show 1000 Ways to Die was originally broadcast in which network?", "The CW"]
            ],
            examples_per_page=10,
            inputs=[input, para_input, target],
            label='Evaluation Examples'
        )  
    with gr.Row():
        button4gen = gr.Button("Generate",variant="primary")

    with gr.Row():
        gr.HTML(
            """
            <h4>Locality</h4>
            """
        )
    with gr.Row():
        loc_input = gr.Dropdown(
            choices=[
                "nq question: where does the phrase good bye felicia come from",
                "nq question: which best describes timbuktu under the mali empire",
                "nq question: where do the question marks go in spanish",
                "nq question: who replaces the vice president in the senate",
                "nq question: active transport performs which function in a cell"
                ],
            value="nq question: which best describes timbuktu under the mali empire",
            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")
   
    button4edit.click(fn=continuous_edit, inputs=[edit_alg, prompt,target_new, num_steps, edit_lr], outputs=[input, target])
    button4gen.click(fn=continuous_union_generate, inputs=[input, para_input, target, edit_alg], outputs=[button4gen_ori, button4gen_edit, button4gen_para_ori, button4gen_para_edit])
    # button4gen.click(fn=generate, inputs=[para_input, target, edit_alg], outputs=[button4gen_para_ori, button4gen_para_edit])
    button4locgen.click(fn=continuous_generate, inputs=[loc_input, edit_alg], outputs=[button4gen_loc_ori, button4gen_loc_edit])


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>]
                    </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("Explanation", open=False):
        gr.Markdown(
            """
            - `Edit Algorithm`: editing method. Choices: [[WISE](https://arxiv.org/abs/2405.14768), [GRACE](https://arxiv.org/abs/2211.11031), [ROME](https://arxiv.org/abs/2202.05262)]
            - `Edit Steps`: the number of times a layer is trained in the editing method.
            - `Edit LR (learning rate)`: the optimization strategy during fine-tuning.
            - `Reliability Evaluation`: the assessment of whether the target edit can be accomplished.
            - `Generalization Evaluation`: whether generalize to the unseen paraphrase prompt.
            - `Locality Evaluation`: the assessment of whether unrelated content has been affected.
            """
        )

    with gr.Tab("Single Knowledge Editing"):
        single_edit_tab()

    with gr.Tab("Continuous Knowledge Editing"):
        continuous_edit_tab()

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

demo.launch()