add continuous

app.py

@@ -1,13 +1,20 @@
 import gradio as gr
 from utils import *
 from transformers import pipeline
-
-css = """
-
-"""
+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()):
@@ -20,91 +27,199 @@ def slowly_reverse(word, progress=gr.Progress()):
         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>""")
-
+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():
-    #     gr.Markdown("<p align='center'><a href='https://github.com/zjunlp/EasyEdit'>🔧https://github.com/zjunlp/EasyEdit</a></p>") 
-    
+    #     input_text = gr.Textbox(label="Status Information",value="Model editing may take about a minute, please be patient.")
     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(
+        gr.HTML(
             """
-            Edit Steps: the number of times a layer is trained in the GRACE method.
+            <h3>Evaluation</h3>
             """
         )
-        gr.Markdown(
+    with gr.Row():
+        gr.HTML(
             """
-            Replacement: the optimization strategy during fine-tuning.
+            <h4>Reliability</h4>
             """
         )
-        gr.Markdown(
-            """
-            Reliability Evaluation: the optimization strategy during fine-tuning.
-            """
-        )       
-        gr.Markdown(
+    # 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(
             """
-            Reliability Evaluation: the assessment of whether the target edit can be accomplished.
+            <h4>Generalization</h4>
             """
-        )          
-        gr.Markdown(
+        )
+    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(
             """
-            Locality Evaluation: the assessment of whether unrelated content has been affected..
+            <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.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')
-        replacement = gr.Dropdown(
-            choices=["replace_last", "replace_all", "replace_prompt"],
-            value="replace_last",
-            label="Replacement",
+        edit_lr = gr.Dropdown(
+            choices=[0.1, 0.5, 1.0],
+            value=1.0,
+            label="Edit LR (learning rate)",
         )
-    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"]
+                ["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=3,
+            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>Reliability Evaluation</h3>
+            <h3>Evaluation</h3>
             """
         )
     with gr.Row():
-        input = gr.Textbox(label="Input Text")
+        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",
@@ -119,25 +234,65 @@ with gr.Blocks(css=css,theme=gr.themes.Soft(text_size="sm")) as demo:
                 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(
             """
-            <h3>Locality Evaluation</h3>
+            <h4>Locality</h4>
             """
         )
     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"
+                "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="where is the ucla usc game being played",
+            value="nq question: which best describes timbuktu under the mali empire",
             label="Unrelated Input Text",
         )
     with gr.Row():
@@ -158,20 +313,76 @@ with gr.Blocks(css=css,theme=gr.themes.Soft(text_size="sm")) as demo:
     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])
+    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 Steps: the number of times a layer is trained in the editing method.
+            """
+        )
+        gr.Markdown(
+            """
+            Edit LR (learning rate): 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.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{wang2023easyedit,
+            @misc{wang2024easyedit,
                 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},
+                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}
@@ -180,5 +391,4 @@ with gr.Blocks(css=css,theme=gr.themes.Soft(text_size="sm")) as demo:
             """
         )
 
-
 demo.launch()

easyeditor/models/__init__.py

@@ -1 +1,3 @@
 from .grace import *
+from .wise import *
+from .rome import *

easyeditor/models/grace/GRACE.py

@@ -29,13 +29,13 @@
 #         layer = config.inner_params[0]
 #         self.device = device
 
-#         # --- ensure proper formatting (GRACE edits ~layers~ not weights matrices) ---        
+#         # --- ensure proper formatting (GRACE edits ~layers~ not weights matrices) ---
 #         suffixes = [".weight", ".bias"]
 #         self.layer = layer.rsplit(".", 1)[0] if any(layer.endswith(x) for x in suffixes) else layer
-        
+
 #         for n, p in self.model.named_parameters():
 #             p.requires_grad = False
-        
+
 #         if isinstance(self.model, transformers.models.gpt2.modeling_gpt2.GPT2LMHeadModel):
 #             transpose = False
 #         else:
@@ -48,32 +48,32 @@
 
 #         if type(original_layer) is not GRACEAdapter:
 #             setattr(edit_module, layer_name, GRACEAdapter(config, original_layer, transpose=transpose).to(self.device))
-        
+
 #     def __call__(self, **kwargs):
 #         # if self.config.task == "hallucination":
 #         #     print(kwargs)
 #         #     key_id = (kwargs["labels"] == -100).sum() - 1
 #         #     setattr(eval(f"self.model.{self.layer}"), "key_id", key_id) # Tell GRACE which token to use for its query (default is the last token)
 #         return self.model(**kwargs)
-    
+
 #     def generate(self, *args, **kwargs):
 #         setattr(eval(f"self.model.{self.layer}"), "key_id", -1)
 #         return self.model.generate(*args, **kwargs)
-        
+
 #     def edit(self, config, tokens):
 #         key_id = (tokens["labels"] == -100).sum() - 1
 #         setattr(eval(f"self.model.{self.layer}"), "key_id", key_id)
-        
+
 #         # --- pass edit label, training mode, and key_id into GRACE ---
 #         setattr(eval(f"self.model.{self.layer}"), "training", True)
 #         setattr(eval(f"self.model.{self.layer}"), "edit_label", tokens["labels"])
-                
+
 #         self.losses = []
 #         # --- train GRACE value ---
 #         for i in range(config.n_iter):
 #             # --- insert iteration into each layer (only initiate keys on iteration 1) ---
 #             setattr(eval(f"self.model.{self.layer}"), "iter", i)
-            
+
 #             # --- pass tokens through model (including through the GRACE layer) ---
 #             outputs = self.model(**tokens)
 #             if i == 0:
@@ -84,14 +84,14 @@
 #             optimizer.step()
 #             optimizer.zero_grad()
 #             self.losses.append(loss.detach().cpu().numpy())
-        
+
 #         self.loss = loss # Log final loss
 
 #         # --- pull out info we want to log from the GRACE layer ---
 #         setattr(eval(f"self.model.{self.layer}"), "training", False)
 #         chosen_key = getattr(eval(f"self.model.{self.layer}"), "chosen_key")
 #         nkeys = len(getattr(eval(f"self.model.{self.layer}"), "keys"))
-            
+
 #         self.log_dict["chosen_key"] =  chosen_key
 #         self.log_dict["nkeys"] = nkeys
 
@@ -109,7 +109,7 @@
 #         self.num_pert = config.num_pert
 #         self.key_id = -1
 #         self.ensure_replace_token_loc = False
-    
+
 #         if transpose:
 #             self.key_shape = layer.weight.shape[1]
 #             self.value_shape = layer.weight.shape[0]
@@ -142,7 +142,7 @@
 #     def split_epsilons_in_half(self, nearest_key, smallest_distance):
 #         self.epsilons[nearest_key] = (smallest_distance / 2) - 1e-5 # Cut nearest epsilon in half
 #         self.epsilons[-1] = smallest_distance / 2 # Cut new epsilon in half
-    
+
 #     def forward(self, *args):
 #         # Run layer forward and save what it would have returned for this instance
 #         layer_out = self.layer(*args)
@@ -176,7 +176,7 @@
 #                 smallest_distance, nearest_key = dists.min(0)
 
 #                 if smallest_distance > (self.init_epsilon + self.epsilons[nearest_key]):
-#                     # If there's no close key, make a new key                    
+#                     # If there's no close key, make a new key
 #                     self.keys, self.values, self.epsilons, self.key_labels = self.add_key(query, new_value)
 #                 else:
 #                     # If there is a close key, we need to handle conflicts
@@ -222,23 +222,27 @@ import torch
 from .utils import parent_module, brackets_to_periods
 import transformers
 import os
+
 os.environ['CUDA_LAUNCH_BLOCKING'] = "1"
 
+
 def euc(query, key):
     # Euclidean distance
     if len(key.shape) < 2:
         key = key.view(1, -1)
     return torch.cdist(key, query, p=2)
 
+
 def perturb_values(chosen_value, num_pert, device):
     # Create a bunch of noised versions of the value, then create batch, then train value
     chosen_value = chosen_value
     noise = torch.normal(0, 1, chosen_value.shape, device=device)
-    noise[0] = noise[0]*0
+    noise[0] = noise[0] * 0
     noise.requires_grad = True
     chosen_value = chosen_value + noise
     return chosen_value
 
+
 class GRACE(torch.nn.Module):
     def __init__(self, config, model, device):
         super(GRACE, self).__init__()
@@ -251,26 +255,27 @@ class GRACE(torch.nn.Module):
         self.device = device
         self.original_layer = None
 
-        # --- ensure proper formatting (GRACE edits ~layers~ not weights matrices) ---        
+        # --- ensure proper formatting (GRACE edits ~layers~ not weights matrices) ---
         suffixes = [".weight", ".bias"]
         self.layer = layer.rsplit(".", 1)[0] if any(layer.endswith(x) for x in suffixes) else layer
-        
+
         for n, p in self.model.named_parameters():
             p.requires_grad = False
-        
+
         if isinstance(self.model, transformers.models.gpt2.modeling_gpt2.GPT2LMHeadModel):
             transpose = False
         else:
             transpose = True
 
         # --- Add GRACE to chosen layers ---
-        edit_module = parent_module(self.model, brackets_to_periods(self.layer))
-        layer_name = self.layer.rsplit(".", 1)[-1]
-        original_layer = getattr(edit_module, layer_name)
+        self.edit_module = parent_module(self.model, brackets_to_periods(self.layer))
+        self.layer_name = self.layer.rsplit(".", 1)[-1]
+        original_layer = getattr(self.edit_module, self.layer_name)
         if type(original_layer) is not GRACEAdapter:
-            setattr(edit_module, layer_name, GRACEAdapter(config, original_layer, transpose=transpose).to(self.device))
+            setattr(self.edit_module, self.layer_name,
+                    GRACEAdapter(config, original_layer, transpose=transpose).to(self.device))
             self.original_layer = copy.deepcopy(original_layer)
-        
+
     def __call__(self, **kwargs):
         # if self.config.task == "hallucination":
         #     print(kwargs)
@@ -278,55 +283,65 @@ class GRACE(torch.nn.Module):
         #     setattr(eval(f"self.model.{self.layer}"), "key_id", key_id) # Tell GRACE which token to use for its query (default is the last token)
         return self.model(**kwargs)
 
+    def get_adapter_layer(self):
+        adapter_layer = getattr(self.edit_module, self.layer_name)
+        assert type(adapter_layer) is GRACEAdapter, print('Adapter Layer is not added correctly....')
+        return adapter_layer
+
     def reset_layer(self):
-        layer_name = self.layer.rsplit(".", 1)[-1]
-        edit_module = parent_module(self.model, brackets_to_periods(self.layer))
-        setattr(edit_module, layer_name, self.original_layer.to(self.device))
+        layer = getattr(self.edit_module, self.layer_name)
+        del layer
+        setattr(self.edit_module, self.layer_name, self.get_adapter_layer().original_layer)
 
     def generate(self, *args, **kwargs):
         setattr(eval(f"self.model.{self.layer}"), "key_id", -1)
         return self.model.generate(*args, **kwargs)
-        
+
     def edit(self, config, tokens):
         key_id = (tokens["labels"] == -100).sum() - 1
         setattr(eval(f"self.model.{self.layer}"), "key_id", key_id)
-        
+
         # --- pass edit label, training mode, and key_id into GRACE ---
         setattr(eval(f"self.model.{self.layer}"), "training", True)
         setattr(eval(f"self.model.{self.layer}"), "edit_label", tokens["labels"])
-                
+
         self.losses = []
         # --- train GRACE value ---
         for i in range(config.n_iter):
             # --- insert iteration into each layer (only initiate keys on iteration 1) ---
             setattr(eval(f"self.model.{self.layer}"), "iter", i)
-            
+
             # --- pass tokens through model (including through the GRACE layer) ---
             outputs = self.model(**tokens)
             if i == 0:
                 # --- we only need to create an optimizer for the first iteration (but forward pass instantiates the key, so optimzer is passed after first inference) ---
                 optimizer = torch.optim.Adam(self.model.parameters(), config.edit_lr)
             loss = outputs.loss
-            loss.backward()
-            optimizer.step()
-            optimizer.zero_grad()
-            self.losses.append(loss.detach().cpu().numpy())
-        
-        self.loss = loss # Log final loss
+            try:
+                loss.backward()
+                optimizer.step()
+                optimizer.zero_grad()
+                self.losses.append(loss.detach().cpu().numpy())
+            except Exception as e:
+                pass
+
+        self.loss = loss  # Log final loss
 
         # --- pull out info we want to log from the GRACE layer ---
         setattr(eval(f"self.model.{self.layer}"), "training", False)
         chosen_key = getattr(eval(f"self.model.{self.layer}"), "chosen_key")
         nkeys = len(getattr(eval(f"self.model.{self.layer}"), "keys"))
-            
-        self.log_dict["chosen_key"] =  chosen_key
+
+        self.log_dict["chosen_key"] = chosen_key
         self.log_dict["nkeys"] = nkeys
 
+
 class GRACEAdapter(torch.nn.Module):
     def __init__(self, config, layer, transpose):
         super(GRACEAdapter, self).__init__()
 
         self.layer = layer
+        self.original_layer = copy.deepcopy(self.layer)
         self.weight = self.layer.weight
         self.init_epsilon = config.eps
         self.dist_fn = config.dist_fn
@@ -335,8 +350,7 @@ class GRACEAdapter(torch.nn.Module):
         self.config = config
         self.num_pert = config.num_pert
         self.key_id = -1
-        self.ensure_replace_token_loc = False
-    
+
         if transpose:
             self.key_shape = layer.weight.shape[1]
             self.value_shape = layer.weight.shape[0]
@@ -346,14 +360,15 @@ class GRACEAdapter(torch.nn.Module):
         self.training = False
 
     def add_key(self, new_key, new_value):
-        keys = torch.vstack([self.keys, new_key.detach()]) # Add new key to list of keys
+        keys = torch.vstack([self.keys, new_key.detach()])  # Add new key to list of keys
 
-        values = torch.nn.Parameter(torch.vstack([self.values, new_value]), requires_grad=True) # Add new value to list of values
+        values = torch.nn.Parameter(torch.vstack([self.values, new_value]),
+                                    requires_grad=True)  # Add new value to list of values
 
         new_epsilon = torch.tensor(self.init_epsilon, device=self.device).view(1)
-        epsilons = torch.vstack([self.epsilons, new_epsilon]) # Add new epsilon to list of epsilons
+        epsilons = torch.vstack([self.epsilons, new_epsilon])  # Add new epsilon to list of epsilons
 
-        key_labels = self.key_labels + [self.edit_label] # Add new key_label to list of key_labels
+        key_labels = self.key_labels + [self.edit_label]  # Add new key_label to list of key_labels
 
         return keys, values, epsilons, key_labels
 
@@ -367,9 +382,9 @@ class GRACEAdapter(torch.nn.Module):
         return edit_label.float().mean() == key_label.float().mean()
 
     def split_epsilons_in_half(self, nearest_key, smallest_distance):
-        self.epsilons[nearest_key] = (smallest_distance / 2) - 1e-5 # Cut nearest epsilon in half
-        self.epsilons[-1] = smallest_distance / 2 # Cut new epsilon in half
-    
+        self.epsilons[nearest_key] = (smallest_distance / 2) - 1e-5  # Cut nearest epsilon in half
+        self.epsilons[-1] = smallest_distance / 2  # Cut new epsilon in half
+
     def forward(self, *args):
         # Run layer forward and save what it would have returned for this instance
         layer_out = self.layer(*args)
@@ -380,13 +395,15 @@ class GRACEAdapter(torch.nn.Module):
             # print(self.__dict__)
             return layer_out
         else:
-            if not self.training and not self.ensure_replace_token_loc and self.key_id == -1:
-                token_to_edit = args[0].shape[1]-1
-                self.key_id = args[0].shape[1]-1
-                self.ensure_replace_token_loc = True
+            if not self.training:
+                if self.key_id == -1:
+                    token_to_edit = args[0].shape[1] - 1
+                    self.key_id = args[0].shape[1] - 1
+                else:
+                    token_to_edit = min(self.key_id, args[0].shape[1] - 1)
             else:
-                token_to_edit = min(self.key_id, args[0].shape[1]-1) # args[0].shape[1] - 1 is sequence length
-            query = args[0][:, token_to_edit, :] # Just use activation for last token
+                token_to_edit = min(self.key_id, args[0].shape[1] - 1)  # args[0].shape[1] - 1 is sequence length
+            query = args[0][:, token_to_edit, :]  # Just use activation for last token
             if self.config.val_init == "cold":
                 new_value = torch.nn.Parameter(torch.rand(1, self.value_shape, requires_grad=True, device=self.device))
             elif self.config.val_init == "warm":
@@ -403,7 +420,7 @@ class GRACEAdapter(torch.nn.Module):
                 smallest_distance, nearest_key = dists.min(0)
 
                 if smallest_distance > (self.init_epsilon + self.epsilons[nearest_key]):
-                    # If there's no close key, make a new key                    
+                    # If there's no close key, make a new key
                     self.keys, self.values, self.epsilons, self.key_labels = self.add_key(query, new_value)
                 else:
                     # If there is a close key, we need to handle conflicts
@@ -413,11 +430,13 @@ class GRACEAdapter(torch.nn.Module):
                     else:
                         # If the current label is the SAME as the nearest label, just make the nearest epsilon bigger
                         if smallest_distance > self.epsilons[nearest_key]:
-                            if self.config.eps_expand== "coverage":
-                                self.epsilons[nearest_key] = smallest_distance # Replace nearest epsilon with dist between old key and new key
+                            if self.config.eps_expand == "coverage":
+                                self.epsilons[
+                                    nearest_key] = smallest_distance  # Replace nearest epsilon with dist between old key and new key
                             elif self.config.eps_expand == "moving_average":
                                 a = 0.5
-                                self.keys[nearest_key] = a*self.keys[nearest_key] + (1-a)*query # Move old key to be halfway between
+                                self.keys[nearest_key] = a * self.keys[nearest_key] + (
+                                            1 - a) * query  # Move old key to be halfway between
                                 self.epsilons[nearest_key] = smallest_distance
                                 # self.epsilons[nearest_key] = smallest_distance + self.init_epsilon
             else:
@@ -435,11 +454,13 @@ class GRACEAdapter(torch.nn.Module):
             chosen_value = perturb_values(chosen_value, self.num_pert, self.device)
 
         if self.replacement == "replace_all":
-            layer_out = torch.where((smallest_dist <= eps).view(-1, 1, 1), chosen_value.unsqueeze(1).repeat_interleave(layer_out.shape[1], 1), layer_out)
+            layer_out = torch.where((smallest_dist <= eps).view(-1, 1, 1),
+                                    chosen_value.unsqueeze(1).repeat_interleave(layer_out.shape[1], 1), layer_out)
         elif self.replacement == "replace_last":
             layer_out[:, token_to_edit] = torch.where((smallest_dist <= eps), chosen_value, layer_out[:, token_to_edit])
         elif self.replacement == "replace_prompt":
-            layer_out[:, :token_to_edit] = torch.where((smallest_dist <= eps), chosen_value, layer_out[:, :token_to_edit])
+            layer_out[:, :token_to_edit] = torch.where((smallest_dist <= eps), chosen_value,
+                                                       layer_out[:, :token_to_edit])
         else:
             print("token replacement choice not found")
         return layer_out

easyeditor/models/grace/grace_main.py

@@ -15,7 +15,7 @@ def apply_grace_to_model(
         requests: List[Dict],
         hparams: GraceHyperParams,
         num_steps: int,
-        replacement: str,
+        edit_lr: float,
         copy=False,
         return_orig_weights=False,
         keep_original_weight=False,
@@ -26,14 +26,13 @@ def apply_grace_to_model(
         model = deepcopy(model)
     weights_copy = {}
     device = torch.device('cpu')
-    hparams.n_iter = num_steps
-    hparams.replacement = replacement
+    hparams.edit_lr = edit_lr
     editor = GRACE(model=model, config=hparams, device=device)
 
     tokens = tokenize(request, tokenizer=tok, device=device)
     editor.edit(config=hparams, tokens=tokens)
 
-    editor.to('cpu')
+    # editor.to('cpu')
     gr.Info("Completed editing via GRACE!")
     return editor
 

easyeditor/models/rome/README.md

@@ -0,0 +1,12 @@
+# ROME
+This package provides a self-contained implementation of Rank-One Model Editing (ROME).
+
+Recall that ROME's update consists of: $u$ selection, $v_*$ optimization, and $v$ insertion.
+* [`compute_u.py`](compute_u.py): Chooses a $u$ vector.
+* [`compute_v.py`](compute_v.py): Choose a $v_*$ via optimization, then computes $v$.
+* [`rome_main.py`](rome_main.py): Instruments main logic.
+* [`rome_params.py`](rome_hparams.py): Interface for specifying hyperparameters. Inherits from the base [`params.py`](../util/hparams.py) module.
+
+For estimating second moment statistics of keys ($C = KK$), we provide the `layer_stats` module. See the [main README](../README.md) for usage instructions.
+* [`layer_stats.py`](layer_stats.py): Logic for retrieving and caching key statistics.
+* [`tok_dataset.py`](tok_dataset.py): Utilities for creating a dataset of tokens.

easyeditor/models/rome/__init__.py

@@ -0,0 +1 @@
+from .rome_main import ROMEHyperParams, apply_rome_to_model, execute_rome

easyeditor/models/rome/compute_u.py

@@ -0,0 +1,125 @@
+import os
+from pathlib import Path
+from typing import Dict, List
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from ..rome import repr_tools
+from ...util.globals import *
+
+from .layer_stats import layer_stats
+from .rome_hparams import ROMEHyperParams
+
+# Cache variables
+inv_mom2_cache = {}
+
+
+def get_inv_cov(
+    model: AutoModelForCausalLM,
+    tok: AutoTokenizer,
+    layer_name: str,
+    mom2_dataset: str,
+    mom2_n_samples: str,
+    mom2_dtype: str,
+    hparams=None,
+) -> torch.Tensor:
+    """
+    Retrieves covariance statistics, then computes the algebraic inverse.
+    Caches result for future use.
+    """
+
+    global inv_mom2_cache
+
+    model_name = model.config._name_or_path.replace("/", "_")
+    key = (model_name, layer_name)
+
+    if key not in inv_mom2_cache:
+        print(
+            f"Retrieving inverse covariance statistics for {model_name} @ {layer_name}. "
+            f"The result will be cached to avoid repetitive computation."
+        )
+        stat = layer_stats(
+            model,
+            tok,
+            layer_name,
+            hparams.stats_dir,
+            mom2_dataset,
+            to_collect=["mom2"],
+            sample_size=mom2_n_samples,
+            precision=mom2_dtype,
+            hparams=hparams
+        )
+        inv_mom2_cache[key] = torch.inverse(
+            stat.mom2.moment().to(f"cuda:{hparams.device}")
+        ).float()  # Cast back to float32
+
+    return inv_mom2_cache[key]
+
+
+def compute_u(
+    model: AutoModelForCausalLM,
+    tok: AutoTokenizer,
+    request: Dict,
+    hparams: ROMEHyperParams,
+    layer: int,
+    context_templates: List[str],
+) -> torch.Tensor:
+    """
+    Computes the right vector used in constructing the rank-1 update matrix.
+    """
+
+    print("Computing left vector (u)...")
+
+    # Compute projection token
+    word_repr_args = dict(
+        model=model,
+        tok=tok,
+        layer=layer,
+        module_template=hparams.rewrite_module_tmp,
+        track="in",
+    )
+    if "subject_" in hparams.fact_token and hparams.fact_token.index("subject_") == 0:
+        word = request["subject"]
+        print(f"Selected u projection object {word}")
+        
+        cur_repr = repr_tools.get_reprs_at_word_tokens(
+            context_templates=[
+                templ.format(request["prompt"]) for templ in context_templates
+            ],
+            words=[word for _ in range(len(context_templates))],
+            subtoken=hparams.fact_token[len("subject_") :],
+            **word_repr_args,
+        ).mean(0)
+
+    elif hparams.fact_token == "last":
+        # Heuristic to choose last word. Not a huge deal if there's a minor
+        # edge case (e.g. multi-token word) because the function below will
+        # take the last token.
+        cur_repr = repr_tools.get_reprs_at_idxs(
+            contexts=[
+                templ.format(request["prompt"].format(request["subject"]))
+                for templ in context_templates
+            ],
+            idxs=[[-1] for _ in range(len(context_templates))],
+            **word_repr_args,
+        ).mean(0)
+        print("Selected u projection token with last token")
+    else:
+        raise ValueError(f"fact_token={hparams.fact_token} not recognized")
+
+    # Apply inverse second moment adjustment
+    u = cur_repr
+    if hparams.mom2_adjustment:
+        u = get_inv_cov(
+            model,
+            tok,
+            hparams.rewrite_module_tmp.format(layer),
+            hparams.mom2_dataset,
+            hparams.mom2_n_samples,
+            hparams.mom2_dtype,
+            hparams=hparams,
+        ) @ u.unsqueeze(1)
+        u = u.squeeze()
+
+    return u / u.norm()

easyeditor/models/rome/compute_v.py

@@ -0,0 +1,278 @@
+from typing import Dict, List, Tuple
+
+import numpy as np
+import torch
+from matplotlib.style import context
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from ..rome import repr_tools
+from ...util import nethook
+
+from .rome_hparams import ROMEHyperParams
+
+
+def compute_v(
+    model: AutoModelForCausalLM,
+    tok: AutoTokenizer,
+    request: Dict,
+    hparams: ROMEHyperParams,
+    layer: int,
+    left_vector: torch.Tensor,
+    context_templates: List[str],
+) -> torch.Tensor:
+    """
+    Computes the value (right) vector for the rank-1 update.
+    Runs a simple optimization procedure.
+    """
+
+    print("Computing right vector (v)")
+
+    # Tokenize target into list of int token IDs
+    target_ids = tok.encode(request["target_new"], return_tensors="pt", add_special_tokens=False).to('cpu')[0]
+
+    # if target_ids[0] == tok.bos_token_id or target_ids[0] == tok.unk_token_id:
+    #     target_ids = target_ids[1:]
+    # Compile list of rewriting and KL x/y pairs
+    rewriting_prompts, kl_prompts = [
+        context.format(request["prompt"]) + tok.decode(target_ids[:-1])
+        for context in context_templates
+    ], ["{} is a"]
+    all_prompts = rewriting_prompts + kl_prompts
+
+    input_tok = tok(
+        [prompt.format(request["subject"]) for prompt in all_prompts],
+        return_tensors="pt",
+        padding=True,
+    ).to("cpu")
+
+    # Compute rewriting targets
+    rewriting_targets = torch.tensor(-100, device='cpu').repeat(
+        len(rewriting_prompts), *input_tok["input_ids"].shape[1:]
+    )
+    for i in range(len(rewriting_prompts)):
+        ex_len = input_tok["attention_mask"][i].sum()
+        rewriting_targets[i, ex_len - len(target_ids) : ex_len] = target_ids
+
+    # Compute indices of the tokens where the fact is looked up
+    vanilla_input_prompts = [
+        context.format(request["prompt"]).format(request['subject'])
+        for context in context_templates
+    ] + [f"{request['subject']} is a"]
+    lookup_idxs = [
+        find_fact_lookup_idx(
+            prompt, request["subject"], tok, hparams.fact_token, verbose=(i == 0), input_prompt=vanilla_input_prompts[i]
+        )
+        for i, prompt in enumerate(all_prompts)
+    ]
+
+    # Finalize rewrite and loss layers
+    loss_layer = max(hparams.v_loss_layer, layer)
+    print(f"Rewrite layer is {layer}")
+    print(f"Tying optimization objective to {loss_layer}")
+
+    # Set up an optimization over a latent vector that, when output at the
+    # rewrite layer, i.e. hypothesized fact lookup location, will induce the
+    # target token to be predicted at the final layer.
+    if hasattr(model.config, 'n_embd'):
+        delta = torch.zeros((model.config.n_embd,), requires_grad=True, device=f"cpu")
+    else:
+        delta = torch.zeros((model.config.hidden_size,), requires_grad=True, device=f"cpu")
+    target_init, kl_distr_init = None, None
+
+    # Inserts new "delta" variable at the appropriate part of the computation
+    def edit_output_fn(cur_out, cur_layer):
+        nonlocal target_init
+        if cur_layer == hparams.mlp_module_tmp.format(layer):
+            # Store initial value of the vector of interest
+            if target_init is None:
+                print("Recording initial value of v*")
+                # Initial value is recorded for the clean sentence
+                target_init = cur_out[0, lookup_idxs[0]].detach().clone()
+                
+            for i, idx in enumerate(lookup_idxs):
+                if len(lookup_idxs)!=len(cur_out):
+                    cur_out[idx, i, :] += delta
+                else:
+                    cur_out[i, idx, :] += delta
+
+        return cur_out
+
+    # Optimizer
+    opt = torch.optim.Adam([delta], lr=hparams.v_lr)
+    nethook.set_requires_grad(False, model)
+
+    # Execute optimization
+    for it in range(hparams.v_num_grad_steps):
+        opt.zero_grad()
+
+        # Forward propagation
+        with nethook.TraceDict(
+            module=model,
+            layers=[
+                hparams.layer_module_tmp.format(loss_layer),
+                hparams.mlp_module_tmp.format(layer),
+            ],
+            retain_input=False,
+            retain_output=True,
+            edit_output=edit_output_fn,
+        ) as tr:
+            logits = model(**input_tok).logits
+
+            # Compute distribution for KL divergence
+            kl_logits = torch.stack(
+                [
+                    logits[i - len(kl_prompts), idx, :]
+                    for i, idx in enumerate(lookup_idxs[-len(kl_prompts) :])
+                ],
+                dim=0,
+            )
+            kl_log_probs = torch.nn.functional.log_softmax(kl_logits, dim=1)
+            if kl_distr_init is None:
+                kl_distr_init = kl_log_probs.detach().clone()
+
+        # Compute loss on rewriting targets
+        log_probs = torch.log_softmax(logits, dim=2)
+
+        loss = torch.gather(
+            log_probs,
+            2,
+            torch.where(rewriting_targets != -100, rewriting_targets, 0).unsqueeze(2),
+        ).squeeze(2)
+        mask = (rewriting_targets != -100).float()
+
+        # Aggregate total losses
+        nll_loss_each = -(loss * mask).sum(1) / target_ids.size(0)
+        nll_loss = nll_loss_each.mean()
+        kl_loss = hparams.kl_factor * torch.nn.functional.kl_div(
+            kl_distr_init, kl_log_probs, log_target=True, reduction="batchmean"
+        )
+        weight_decay = hparams.v_weight_decay * (
+            torch.norm(delta) / torch.norm(target_init) ** 2
+        )
+        # weight_decay = hparams.v_weight_decay * torch.norm(delta) ** 2
+        loss = nll_loss + kl_loss + weight_decay
+        print(
+            f"loss {np.round(loss.item(), 3)} = {np.round(nll_loss.item(), 3)} + {np.round(kl_loss.item(), 3)} + {np.round(weight_decay.item(), 3)} "
+            f"avg prob of [{request['target_new']}] "
+            f"{torch.exp(-nll_loss_each).mean().item()}"
+        )
+        if loss < 5e-2:
+            break
+
+        if it == hparams.v_num_grad_steps - 1:
+            break
+
+        # Backpropagate
+        loss.backward()
+        opt.step()
+
+        # Project within L2 ball
+        max_norm = hparams.clamp_norm_factor * target_init.norm()
+        if delta.norm() > max_norm:
+            with torch.no_grad():
+                delta[...] = delta * max_norm / delta.norm()
+
+    target = target_init + delta.to(target_init.dtype)
+
+    # Retrieve cur_input, the current input to the 2nd MLP layer, and
+    # cur_output, the original output of the 2nd MLP layer.
+    cur_input, cur_output = get_module_input_output_at_word(
+        model,
+        tok,
+        layer,
+        context_template=request["prompt"],
+        word=request["subject"],
+        module_template=hparams.rewrite_module_tmp,
+        fact_token_strategy=hparams.fact_token,
+    )
+
+    # Solving the linear system to compute the right vector
+    right_vector = (target - cur_output) / torch.dot(cur_input, left_vector)
+    print(f"Delta norm: {(target - cur_output).norm().item()}")
+    print(
+        f"Change in target norm: {target_init.norm().item()} to {target.norm().item()} => {(target.norm() - target_init.norm()).item()}"
+    )
+    print(f"Division Factor: {torch.dot(cur_input, left_vector).item()}")
+    print(f"Right vector norm: {right_vector.norm()}")
+
+    return right_vector
+
+
+def get_module_input_output_at_word(
+    model: AutoModelForCausalLM,
+    tok: AutoTokenizer,
+    layer: int,
+    context_template: str,
+    word: str,
+    module_template: str,
+    fact_token_strategy: str,
+) -> Tuple[torch.Tensor]:
+    """
+    Retrieves detached representations for a word at the input and
+    output of a particular layer module.
+    """
+
+    word_repr_args = dict(
+        model=model,
+        tok=tok,
+        layer=layer,
+        module_template=module_template,
+    )
+    if "subject_" in fact_token_strategy and fact_token_strategy.index("subject_") == 0:
+        subtoken = fact_token_strategy[len("subject_") :]
+        l_input, l_output = repr_tools.get_reprs_at_word_tokens(
+            track="both",
+            subtoken=subtoken,
+            context_templates=[context_template],
+            words=[word],
+            **word_repr_args,
+        )
+    elif fact_token_strategy == "last":
+        l_input, l_output = repr_tools.get_reprs_at_idxs(
+            track="both",
+            contexts=[context_template.format(word)],
+            idxs=[[-1]],
+            **word_repr_args,
+        )
+    else:
+        raise ValueError(f"fact_token={fact_token_strategy} not recognized")
+
+    l_input, l_output = l_input[0], l_output[0]
+    return l_input.detach(), l_output.detach()
+
+
+def find_fact_lookup_idx(
+    prompt: str,
+    subject: str,
+    tok: AutoTokenizer,
+    fact_token_strategy: str,
+    verbose=True,
+    input_prompt=None
+) -> int:
+    """
+    Computes hypothesized fact lookup index given a sentence and subject.
+    """
+
+    ret = None
+    if fact_token_strategy == "last":
+        ret = len(tok.encode(input_prompt)) - 1
+    elif (
+        "subject_" in fact_token_strategy and fact_token_strategy.index("subject_") == 0
+    ):
+        ret = repr_tools.get_words_idxs_in_templates(
+            tok=tok,
+            context_templates=[prompt],
+            words=[subject],
+            subtoken=fact_token_strategy[len("subject_") :],
+        )[0][0]
+    else:
+        raise ValueError(f"fact_token={fact_token_strategy} not recognized")
+
+    sentence = prompt.format(subject)
+    if verbose:
+        print(
+            f"Lookup index found: {ret} | Sentence: {sentence} | Token:",
+            tok.decode(tok(sentence)["input_ids"][ret]),
+        )
+
+    return ret

easyeditor/models/rome/layer_stats.py

@@ -0,0 +1,198 @@
+import os
+from pathlib import Path
+
+import torch
+from datasets import load_dataset
+from tqdm.auto import tqdm
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from ...util.globals import *
+from ...util.nethook import Trace, set_requires_grad
+from ...util.runningstats import CombinedStat, Mean, NormMean, SecondMoment, tally
+
+from .tok_dataset import (
+    TokenizedDataset,
+    dict_to_,
+    flatten_masked_batch,
+    length_collation,
+)
+
+STAT_TYPES = {
+    "mom2": SecondMoment,
+    "mean": Mean,
+    "norm_mean": NormMean,
+}
+
+
+def main():
+    """
+    Command-line utility to precompute cached stats.
+    """
+    import argparse
+
+    parser = argparse.ArgumentParser(description="ROME Statistics Collector")
+
+    def aa(*args, **kwargs):
+        parser.add_argument(*args, **kwargs)
+
+    aa("--model_name", default="gpt2-xl", choices=["gpt2-xl", "EleutherAI/gpt-j-6B"])
+    aa("--dataset", default="wikipedia", choices=["wikitext", "wikipedia"])
+    aa("--layers", default=[17], type=lambda x: list(map(int, x.split(","))))
+    aa("--to_collect", default=["mom2"], type=lambda x: x.split(","))
+    aa("--sample_size", default=100000, type=lambda x: None if x == "all" else int(x))
+    aa("--batch_tokens", default=None, type=lambda x: None if x == "any" else int(x))
+    aa("--precision", default="float32", choices=["float64", "float32", "float16"])
+    aa("--stats_dir", default=STATS_DIR)
+    aa("--download", default=1, type=int, choices=[0, 1])
+    args = parser.parse_args()
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+    model = AutoModelForCausalLM.from_pretrained(args.model_name).eval().cuda()
+    set_requires_grad(False, model)
+
+    for layer_num in args.layers:
+        print(
+            f"Computing stats for layer {layer_num} of {args.model_name} "
+            f'over {args.sample_size or "all"} samples of {args.dataset}. '
+            "Note, the statistics are collected over the inputs to the second MLP layer, "
+            "or equivalently the outputs of the first MLP layer."
+        )
+        proj_layer_name = "c_proj" if "gpt2" in args.model_name else "fc_out"
+        layer_name = f"transformer.h.{layer_num}.mlp.{proj_layer_name}"
+
+        layer_stats(
+            model,
+            tokenizer,
+            layer_name,
+            args.stats_dir,
+            args.dataset,
+            args.to_collect,
+            sample_size=args.sample_size,
+            precision=args.precision,
+            batch_tokens=args.batch_tokens,
+            download=args.download,
+        )
+
+
+def layer_stats(
+    model,
+    tokenizer,
+    layer_name,
+    stats_dir,
+    ds_name,
+    to_collect,
+    model_name=None,
+    sample_size=None,
+    precision=None,
+    batch_tokens=None,
+    download=True,
+    progress=tqdm,
+    force_recompute=False,
+    hparams=None
+):
+    """
+    Function to load or compute cached stats.
+    """
+
+    def get_ds():
+        # Load_From_File
+        # from datasets import Dataset
+        # raw_ds = Dataset.from_file('XXX/XXX/wikipedia-train.arrow')
+        # raw_ds = {'train': raw_ds}
+        raw_ds = load_dataset(
+            ds_name,
+            dict(wikitext="wikitext-103-raw-v1", wikipedia="20200501.en")[ds_name]
+        )
+        if hasattr(model.config, 'n_positions'):
+            maxlen = model.config.n_positions
+        elif hasattr(model.config, 'max_sequence_length'):
+            maxlen = model.config.max_sequence_length
+        elif hasattr(model.config, 'max_position_embeddings'):
+            maxlen = model.config.max_position_embeddings
+        elif hasattr(model.config,'seq_length'):
+            maxlen = model.config.seq_length
+        else:
+            raise NotImplementedError
+                
+        if hasattr(model.config, 'model_type') and 'mistral' in model.config.model_type:
+            if hasattr(model.config, 'sliding_window') and model.config.sliding_window:
+                maxlen = model.config.sliding_window or 4096
+            else:
+                maxlen = 4096
+        
+        if batch_tokens is not None and batch_tokens < maxlen:
+            maxlen = batch_tokens
+        return TokenizedDataset(raw_ds["train"], tokenizer, maxlen=maxlen)
+
+    # Continue with computation of statistics
+    batch_size = 100  # Examine this many dataset texts at once
+    if hasattr(model.config, 'n_positions'):
+        npos = model.config.n_positions
+    elif hasattr(model.config, 'max_sequence_length'):
+        npos = model.config.max_sequence_length
+    elif hasattr(model.config, 'max_position_embeddings'):
+        npos = model.config.max_position_embeddings
+    elif hasattr(model.config,'seq_length'):
+        npos = model.config.seq_length
+    else:
+        raise NotImplementedError
+        
+    if hasattr(model.config, 'model_type') and 'mistral' in model.config.model_type:
+        if hasattr(model.config, 'sliding_window') and model.config.sliding_window:
+            npos = model.config.sliding_window or 4096
+        else:
+            npos = 4096
+    
+    if batch_tokens is None:
+        batch_tokens = npos * 3  # Sort and divide into batches with this many tokens
+    if precision is None:
+        precision = "float64"
+    dtype = getattr(torch, precision)
+    size_suffix = "" if sample_size is None else f"_{sample_size}"
+    if batch_tokens < npos:
+        size_suffix = "_t{batch_tokens}" + size_suffix
+    if model_name is None:
+        # model_name = model.config._name_or_path.replace("/", "_")
+        model_name = model.config._name_or_path.rsplit("/")[-1]
+
+    stats_dir = Path(stats_dir)
+    file_extension = f"{model_name}/{ds_name}_stats/{layer_name}_{precision}_{'-'.join(sorted(to_collect))}{size_suffix}.npz"
+    filename = stats_dir / file_extension
+
+    print(f"Computing Cov locally....")
+
+    ds = get_ds() if not filename.exists() else None
+
+    if progress is None:
+        progress = lambda x: x
+
+    stat = CombinedStat(**{k: STAT_TYPES[k]() for k in to_collect})
+    loader = tally(
+        stat,
+        ds,
+        cache=(filename if not force_recompute else None),
+        sample_size=sample_size,
+        batch_size=batch_size,
+        collate_fn=length_collation(batch_tokens),
+        pin_memory=True,
+        random_sample=1,
+        num_workers=2,
+    )
+    batch_count = -(-(sample_size or len(ds)) // batch_size)
+    with torch.no_grad():
+        for batch_group in progress(loader, total=batch_count):
+            for batch in batch_group:
+                batch = dict_to_(batch, f"cuda:{hparams.device}")
+                with Trace(
+                    model, layer_name, retain_input=True, retain_output=False, stop=True
+                ) as tr:
+                    model(**batch)
+                feats = flatten_masked_batch(tr.input, batch["attention_mask"])
+                # feats = flatten_masked_batch(tr.output, batch["attention_mask"])
+                feats = feats.to(dtype=dtype)
+                stat.add(feats)
+    return stat
+
+
+if __name__ == "__main__":
+    main()

easyeditor/models/rome/repr_tools.py

@@ -0,0 +1,174 @@
+"""
+Contains utilities for extracting token representations and indices
+from string templates. Used in computing the left and right vectors for ROME.
+"""
+
+from copy import deepcopy
+from typing import List
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from ...util import nethook
+
+def get_reprs_at_word_tokens(
+    model: AutoModelForCausalLM,
+    tok: AutoTokenizer,
+    context_templates: List[str],
+    words: List[str],
+    layer: int,
+    module_template: str,
+    subtoken: str,
+    track: str = "in",
+) -> torch.Tensor:
+    """
+    Retrieves the last token representation of `word` in `context_template`
+    when `word` is substituted into `context_template`. See `get_last_word_idx_in_template`
+    for more details.
+    """
+
+    idxs = get_words_idxs_in_templates(tok, context_templates, words, subtoken)
+    return get_reprs_at_idxs(
+        model,
+        tok,
+        [context_templates[i].format(words[i]) for i in range(len(words))],
+        idxs,
+        layer,
+        module_template,
+        track,
+    )
+
+def get_words_idxs_in_templates(
+    tok: AutoTokenizer, context_templates: str, words: str, subtoken: str
+) -> int:
+    """
+    Given list of template strings, each with *one* format specifier
+    (e.g. "{} plays basketball"), and words to be substituted into the
+    template, computes the post-tokenization index of their last tokens.
+    """
+
+    assert all(
+        tmp.count("{}") == 1 for tmp in context_templates
+    ), "We currently do not support multiple fill-ins for context"
+
+
+    prefixes_len, words_len, suffixes_len, inputs_len = [], [], [], []
+    for i, context in enumerate(context_templates):
+        prefix, suffix = context.split("{}")
+        prefix_len = len(tok.encode(prefix))
+        prompt_len = len(tok.encode(prefix + words[i]))
+        input_len = len(tok.encode(prefix + words[i] + suffix))
+        prefixes_len.append(prefix_len)
+        words_len.append(prompt_len - prefix_len)
+        suffixes_len.append(input_len - prompt_len)
+        inputs_len.append(input_len)
+
+    # Compute prefixes and suffixes of the tokenized context
+    # fill_idxs = [tmp.index("{}") for tmp in context_templates]
+    # prefixes, suffixes = [
+    #     tmp[: fill_idxs[i]] for i, tmp in enumerate(context_templates)
+    # ], [tmp[fill_idxs[i] + 2 :] for i, tmp in enumerate(context_templates)]
+    # words = deepcopy(words)
+    #
+    # # Pre-process tokens
+    # for i, prefix in enumerate(prefixes):
+    #     if len(prefix) > 0:
+    #         assert prefix[-1] == " "
+    #         prefix = prefix[:-1]
+    #
+    #         prefixes[i] = prefix
+    #         words[i] = f" {words[i].strip()}"
+    #
+    # # Tokenize to determine lengths
+    # assert len(prefixes) == len(words) == len(suffixes)
+    # n = len(prefixes)
+    # batch_tok = tok([*prefixes, *words, *suffixes])
+    # if 'input_ids' in batch_tok:
+    #     batch_tok = batch_tok['input_ids']
+    # prefixes_tok, words_tok, suffixes_tok = [
+    #     batch_tok[i : i + n] for i in range(0, n * 3, n)
+    # ]
+    # prefixes_len, words_len, suffixes_len = [
+    #     [len(el) for el in tok_list]
+    #     for tok_list in [prefixes_tok, words_tok, suffixes_tok]
+    # ]
+
+    # Compute indices of last tokens
+    if subtoken == "last" or subtoken == "first_after_last":
+        return [
+            [
+                prefixes_len[i]
+                + words_len[i]
+                - (1 if subtoken == "last" or suffixes_len[i] == 0 else 0)
+            ]
+            # If suffix is empty, there is no "first token after the last".
+            # So, just return the last token of the word.
+            for i in range(len(context_templates))
+        ]
+    elif subtoken == "first":
+        return [[prefixes_len[i] - inputs_len[i]] for i in range(len(context_templates))]
+    else:
+        raise ValueError(f"Unknown subtoken type: {subtoken}")
+
+
+def get_reprs_at_idxs(
+    model: AutoModelForCausalLM,
+    tok: AutoTokenizer,
+    contexts: List[str],#表示该知识的完整句子
+    idxs: List[List[int]],#被填入词的位置
+    layer: int,
+    module_template: str,
+    track: str = "in",
+) -> torch.Tensor:
+    """
+    Runs input through model and returns averaged representations of the tokens
+    at each index in `idxs`.
+    """
+
+    def _batch(n):
+        for i in range(0, len(contexts), n):
+            yield contexts[i : i + n], idxs[i : i + n]#将句子和被填词位置分块
+
+    assert track in {"in", "out", "both"}
+    both = track == "both"
+    tin, tout = (
+        (track == "in" or both),
+        (track == "out" or both),
+    )#tin tout都是bool结构
+    module_name = module_template.format(layer)
+    to_return = {"in": [], "out": []}
+
+    def _process(cur_repr, batch_idxs, key):
+        nonlocal to_return
+        cur_repr = cur_repr[0] if type(cur_repr) is tuple else cur_repr
+        if cur_repr.shape[0]!=len(batch_idxs):
+            cur_repr=cur_repr.transpose(0,1)
+        for i, idx_list in enumerate(batch_idxs):
+            to_return[key].append(cur_repr[i][idx_list].mean(0))
+
+    for batch_contexts, batch_idxs in _batch(n=128):
+        #contexts_tok:[21 19]
+        contexts_tok = tok(batch_contexts, padding=True, return_tensors="pt").to(
+            next(model.parameters()).device
+        )
+
+        with torch.no_grad():
+            with nethook.Trace(
+                module=model,
+                layer=module_name,
+                retain_input=tin,
+                retain_output=tout,
+            ) as tr:
+                model(**contexts_tok)
+
+        if tin:
+            _process(tr.input, batch_idxs, "in")
+        if tout:
+            _process(tr.output, batch_idxs, "out")
+
+    to_return = {k: torch.stack(v, 0) for k, v in to_return.items() if len(v) > 0}
+
+    if len(to_return) == 1:
+        return to_return["in"] if tin else to_return["out"]
+    else:
+        return to_return["in"], to_return["out"]

easyeditor/models/rome/rome_hparams.py

@@ -0,0 +1,55 @@
+from dataclasses import dataclass
+from typing import List
+import yaml
+
+from ...util.hparams import HyperParams
+
+
+@dataclass
+class ROMEHyperParams(HyperParams):
+    # Method
+    layers: List[int]
+    fact_token: str
+    v_num_grad_steps: int
+    v_lr: float
+    v_loss_layer: int
+    v_weight_decay: float
+    clamp_norm_factor: float
+    kl_factor: float
+    mom2_adjustment: bool
+    context_template_length_params: List[List[int]]
+
+    # Module templates
+    rewrite_module_tmp: str
+    layer_module_tmp: str
+    mlp_module_tmp: str
+    attn_module_tmp: str
+    ln_f_module: str
+    lm_head_module: str
+
+    # Statistics
+    mom2_dataset: str
+    mom2_n_samples: int
+    mom2_dtype: str
+    alg_name: str
+    device: int
+    model_name: str
+    stats_dir: str
+
+    max_length: int = 40
+    model_parallel: bool = False
+    fp16: bool = False
+
+    @classmethod
+    def from_hparams(cls, hparams_name_or_path: str):
+
+        if '.yaml' not in hparams_name_or_path:
+            hparams_name_or_path = hparams_name_or_path + '.yaml'
+
+        with open(hparams_name_or_path, "r") as stream:
+            config = yaml.safe_load(stream)
+            config = super().construct_float_from_scientific_notation(config)
+
+        assert (config and config['alg_name'] == 'ROME') or print(f'ROMEHyperParams can not load from {hparams_name_or_path}, '
+                                                f'alg_name is {config["alg_name"]} ')
+        return cls(**config)

easyeditor/models/rome/rome_main.py

@@ -0,0 +1,192 @@
+from copy import deepcopy
+from typing import Dict, List, Tuple
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from ...util import nethook
+from ...util.generate import generate_fast
+
+from .compute_u import compute_u
+from .compute_v import compute_v
+from .rome_hparams import ROMEHyperParams
+import gradio as gr
+
+CONTEXT_TEMPLATES_CACHE = None
+
+
+def apply_rome_to_model(
+    model: AutoModelForCausalLM,
+    tok: AutoTokenizer,
+    request: List[Dict],
+    hparams: ROMEHyperParams,
+    num_steps: int,
+    edit_lr: float,
+    copy=False,
+    return_orig_weights=False,
+    keep_original_weight=False,
+    **kwargs
+) -> Tuple[AutoModelForCausalLM, List[str]]:
+    """
+    Returns a model with the desired changes.
+
+    :param copy: If true, will preserve the original model while creating a new one to edit.
+        Note that you are responsible for deallocating the new model's memory to avoid leaks.
+
+    :return: (1) the updated model, (2) an original copy of the weights that changed
+    """
+    if copy:
+        model = deepcopy(model)
+
+    weights_copy = {}
+    hparams.v_num_grad_steps = num_steps // 2
+    hparams.v_lr = edit_lr
+    request['subject'] = request['prompt']
+
+    deltas = execute_rome(model, tok, request, hparams)
+
+    with torch.no_grad():
+        for w_name, (delta_u, delta_v) in deltas.items():
+            upd_matrix = delta_u.unsqueeze(1) @ delta_v.unsqueeze(0)
+            w = nethook.get_parameter(model, w_name)
+            upd_matrix = upd_matrix_match_shape(upd_matrix, w.shape)
+
+            if return_orig_weights and w_name not in weights_copy:
+                weights_copy[w_name] = w.detach().clone()
+
+            w[...] += upd_matrix
+
+        print(f"New weights successfully inserted into {list(deltas.keys())}")
+
+    if not keep_original_weight:
+        weights_copy = {}
+    gr.Info("Completed editing via ROME!")
+    return model
+
+
+def execute_rome(
+    model: AutoModelForCausalLM,
+    tok: AutoTokenizer,
+    request: Dict,
+    hparams: ROMEHyperParams,
+) -> Dict[str, Tuple[torch.Tensor]]:
+    """
+    Executes the ROME update algorithm for the specified update at the specified layer
+    Invariant: model at beginning of function == model at end of function
+    """
+
+    # Update target and print info
+    request = deepcopy(request)
+    if request["target_new"] != " ":
+        # Space required for correct tokenization
+        request["target_new"] = " " + request["target_new"]
+
+    if '{}' not in request['prompt']:
+        assert request['subject'] in request['prompt'] or \
+               print(f"Subject:{request['subject']} do not exist in prompt: {request['prompt']}")
+
+        request['prompt'] = request['prompt'].replace(request['subject'], '{}')
+
+    print(
+        f"Executing ROME algorithm for the update: "
+        f"[{request['prompt'].format(request['subject'])}] -> [{request['target_new']}]"
+    )
+
+    # Retrieve weights that user desires to change
+    weights = {
+        f"{hparams.rewrite_module_tmp.format(layer)}.weight": nethook.get_parameter(
+            model, f"{hparams.rewrite_module_tmp.format(layer)}.weight"
+        )
+        for layer in hparams.layers
+    }
+    # Save old weights for future restoration
+    weights_copy = {k: v.detach().clone() for k, v in weights.items()}
+
+    # Update loop: sequentially intervene at each specified layer
+    deltas = {}
+    for layer in sorted(hparams.layers):
+        # Compute rank-1 update matrix
+        left_vector: torch.Tensor = compute_u(
+            model,
+            tok,
+            request,
+            hparams,
+            layer,
+            get_context_templates(model, tok, hparams.context_template_length_params),
+        )
+        print("Left vector shape:", left_vector.shape)
+        right_vector: torch.Tensor = compute_v(
+            model,
+            tok,
+            request,
+            hparams,
+            layer,
+            left_vector,
+            get_context_templates(model, tok, hparams.context_template_length_params),
+        )
+        print("Right vector shape:", right_vector.shape)
+
+        with torch.no_grad():
+            # Determine correct transposition of delta matrix
+            weight_name = f"{hparams.rewrite_module_tmp.format(layer)}.weight"
+            upd_matrix = left_vector.unsqueeze(1) @ right_vector.unsqueeze(0)
+            upd_matrix = upd_matrix_match_shape(upd_matrix, weights[weight_name].shape)
+
+            # Update model weights and record desired changes in `delta` variable
+            weights[weight_name][...] += upd_matrix
+            deltas[weight_name] = (
+                left_vector.detach(),
+                right_vector.detach(),
+            )
+
+    # Restore state of original model
+    with torch.no_grad():
+        for k, v in weights.items():
+            v[...] = weights_copy[k]
+
+    print(f"Deltas successfully computed for {list(weights.keys())}")
+
+    return deltas
+
+
+def upd_matrix_match_shape(matrix: torch.Tensor, shape: torch.Size) -> torch.Tensor:
+    """
+    GPT-2 and GPT-J have transposed weight representations.
+    Returns a matrix that matches the desired shape, else raises a ValueError
+    """
+
+    if matrix.shape == shape:
+        return matrix
+    elif matrix.T.shape == shape:
+        return matrix.T
+    else:
+        raise ValueError(
+            "Update matrix computed by ROME does not match original weight shape. "
+            "Check for bugs in the code?"
+        )
+
+
+def get_context_templates(model, tok, length_params):
+    global CONTEXT_TEMPLATES_CACHE
+
+    if CONTEXT_TEMPLATES_CACHE is None:
+        CONTEXT_TEMPLATES_CACHE = ["{}"] + [
+            x.replace("{", "").replace("}", "") + ". {}"
+            for x in sum(
+                (
+                    generate_fast(
+                        model,
+                        tok,
+                        ["The", "Therefore", "Because", "I", "You"],
+                        n_gen_per_prompt=n_gen // 5,
+                        max_out_len=length,
+                    )
+                    for length, n_gen in length_params
+                ),
+                [],
+            )
+        ]
+
+        print(f"Cached context templates {CONTEXT_TEMPLATES_CACHE}")
+
+    return CONTEXT_TEMPLATES_CACHE

easyeditor/models/rome/tok_dataset.py

@@ -0,0 +1,99 @@
+import torch
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import Dataset
+
+
+class TokenizedDataset(Dataset):
+    """
+    Converts a dataset of text samples into a dataset of token sequences,
+    as converted by a supplied tokenizer. The tokens come along with position
+    ids and attention masks, they can be supplied direcly to the model.
+    """
+
+    def __init__(self, text_dataset, tokenizer=None, maxlen=None, field="text"):
+        self.text_dataset = text_dataset
+        self.field = field
+        self.tokenizer = tokenizer
+        self.maxlen = maxlen
+        if hasattr(text_dataset, "info"):
+            self.info = text_dataset.info
+
+    def __len__(self):
+        return len(self.text_dataset)
+
+    def __getitem__(self, i):
+        text = self.text_dataset[i]
+        if self.field is not None:
+            text = text[self.field]
+        token_list = self.tokenizer.encode(
+            text, truncation=True, max_length=self.maxlen
+        )
+        position_ids = list(range(len(token_list)))
+        attention_mask = [1] * len(token_list)
+        return dict(
+            input_ids=torch.tensor(token_list),
+            position_ids=torch.tensor(position_ids),
+            attention_mask=torch.tensor(attention_mask),
+        )
+
+
+def dict_to_(data, device):
+    """
+    Moves a dictionary of tensors to the specified device.
+    """
+    for k in data:
+        data[k] = data[k].to(device)
+    return data
+
+
+def length_collation(token_size):
+    """
+    Sorts a batch of sequences and breaks it up into subbatches
+    of same-sized sequences, padding as needed.  Each batch
+    has no more than token_size total tokens (or a single
+    sequence, if the sequence happens to be larger).
+    """
+
+    def collate_fn(items):
+        items = sorted(items, key=lambda x: -len(x["input_ids"]))
+        batches = []
+        batch = []
+        batch_width = 0
+        for item in items:
+            item_width = len(item["input_ids"])
+            if item_width == 0:
+                break
+            if batch_width * (len(batch) + 1) > token_size:
+                batches.append(make_padded_batch(batch))
+                batch = []
+                batch_width = 0
+            if not batch:
+                batch_width = item_width
+            batch.append(item)
+        if len(batch):
+            batches.append(make_padded_batch(batch))
+        return batches
+
+    return collate_fn
+
+
+def make_padded_batch(items):
+    """
+    Pads sequences in a batch, so they are all the same length as the longest.
+    """
+    max_len = max(len(d["input_ids"]) for d in items)
+    if max_len == 0:
+        return {k: torch.zeros((0, 0), dtype=torch.long) for k in items[0]}
+    return {
+        k: pad_sequence([d[k] for d in items if len(d["input_ids"])], batch_first=True)
+        for k, v in items[0].items()
+    }
+
+
+def flatten_masked_batch(data, mask):
+    """
+    Flattens feature data, ignoring items that are masked out of attention.
+    """
+    flat_data = data.view(-1, data.size(-1))
+    attended_tokens = mask.view(-1).nonzero()[:, 0]
+    return flat_data[attended_tokens]

easyeditor/models/wise/WISE.py

@@ -0,0 +1,466 @@
+import copy
+import random
+
+import torch
+from torch.nn import functional as F
+from .utils import parent_module, brackets_to_periods, EarlyStopMeter, EditingMeanAct
+import transformers
+import numpy as np
+from torch import Tensor
+from torch.nn import CrossEntropyLoss
+from transformers.activations import ACT2FN
+from .merge import slerp, GTA, linear
+import torch.nn as nn
+import gc
+
+merge_dict = {
+    'slerp': slerp(),
+    'ties': GTA('magnitude', 'sum', normalize=True),
+    'magnitude_norm': GTA('magnitude', None, normalize=True),
+    'magnitude': GTA('magnitude', None, normalize=False),
+    'sign': GTA(None, 'sum', normalize=True),
+    'dare_ties': GTA('rescaled_random', 'sum'),
+    'dare_linear': GTA('random', None),
+    'linear': linear()
+}
+
+edit_history = []
+merge_group_edit_history = []
+
+def euc(query, key, config, act_mask=None, infer=False):
+    # Euclidean distance
+
+    act_fn = ACT2FN[config.hidden_act]
+    l2_norm = torch.norm(act_fn(key) - act_fn(query), dim=-1)
+    if infer and l2_norm.size(1) > 100:
+        topk = torch.topk(l2_norm, k=1, largest=True)
+        return topk.values.mean()
+
+    if act_mask is not None:
+        return torch.sum(l2_norm * act_mask, dim=1) / torch.sum(act_mask, dim=1)
+    else:
+        return torch.mean(l2_norm, dim=-1)
+
+
+class WISE(torch.nn.Module):
+    def __init__(self, config, model, device):
+        super(WISE, self).__init__()
+        self.config = config
+        self.model = model
+        self.config = config
+        if hasattr(self.model.config, 'hidden_act'):
+            self.config.hidden_act = self.model.config.hidden_act
+        elif hasattr(self.model.config, 'activation_function'):
+            self.config.hidden_act = self.model.config.activation_function
+        # self.tokenizer = model.tokenizer
+        layer = config.inner_params[0]
+        self.device = device
+        self.adapter_layer = None
+        self.original_layer = None
+
+        # --- ensure proper formatting (WISE edits weights matrices) ---
+        suffixes = [".weight", ".bias"]
+        self.layer = layer.rsplit(".", 1)[0] if any(layer.endswith(x) for x in suffixes) else layer
+
+        for n, p in self.model.named_parameters():
+            p.requires_grad = False
+
+        if isinstance(self.model, transformers.models.gpt2.modeling_gpt2.GPT2LMHeadModel):
+            conv1D = True
+        else:
+            conv1D = False
+
+        # --- Add WISE to chosen layers ---
+        self.edit_module = parent_module(self.model, brackets_to_periods(self.layer))
+        self.layer_name = self.layer.rsplit(".", 1)[-1]
+        adapter_layer = getattr(self.edit_module, self.layer_name)
+
+        if type(adapter_layer) is not WISEAdapter:
+            setattr(self.edit_module, self.layer_name, WISEAdapter(config, adapter_layer, conv1D=conv1D))
+            self.original_layer = copy.deepcopy(adapter_layer)
+            print(f"New weights successfully inserted into {layer}")
+        
+        gc.collect()
+        torch.cuda.empty_cache()
+        gc.collect()
+
+    # Forward
+    def __call__(self, **kwargs):
+        if not self.config.retrieve:
+            if hasattr(self.get_adapter_layer(), 'editing') and not self.get_adapter_layer().editing:
+                # final merge
+                if not self.get_adapter_layer().original_layer.weight.equal(self.get_adapter_layer().new_weight) and self.get_adapter_layer().editing_total_cnt >= self.config.save_freq:
+                    self.get_adapter_layer().memory_weight.append(self.get_adapter_layer().new_weight)
+                if len(self.get_adapter_layer().memory_weight) > 0 and self.get_adapter_layer().editing_total_cnt >= self.config.save_freq:
+                    print('length of memory is ', len(self.get_adapter_layer().memory_weight), '!!!!!!')
+                    self.get_adapter_layer().merge_weight()
+        return self.model(**kwargs)
+
+    def reset_layer(self):
+        layer = getattr(self.edit_module, self.layer_name)
+        del layer
+        setattr(self.edit_module, self.layer_name, self.get_adapter_layer().original_layer)
+
+    def get_adapter_layer(self):
+        adapter_layer = getattr(self.edit_module, self.layer_name)
+        assert type(adapter_layer) is WISEAdapter, print('Adapter Layer is not added correctly....')
+        return adapter_layer
+
+    # TODO: generation
+    def generate(self, *args, **kwargs):
+        setattr(eval(f"self.model.{self.layer}"), "key_id", -1)
+        return self.model.generate(*args, **kwargs)
+
+    def edit(self, config, tokens, act_mask=None, deact_mask=None):
+        # for retrieve ##
+        global edit_history
+        global merge_group_edit_history
+        edit_history.append([{f"{k1}" : v1.to('cpu') for k1, v1 in tokens.items()}, False])
+        # for retrieve ##
+        last_prompt_token_loc = (tokens["labels"] == -100).sum(dim=-1) - 1
+
+        setattr(eval(f"self.model.{self.layer}"), "training", True)
+        setattr(eval(f"self.model.{self.layer}"), "editing", True)
+        self.get_adapter_layer().set_parameter_tunable()
+        if getattr(eval(f"self.model.{self.layer}"), "editing_total_cnt") % self.config.save_freq == 0:
+            self.get_adapter_layer().generate_activation_mask(self.config.mask_ratio)
+
+        # --- train Wise value ---
+        loss_meter = EarlyStopMeter()
+        for i in range(config.n_iter):
+
+            if i == 0:
+                # --- we only need to create an optimizer for the first iteration (but forward pass instantiates the key, so optimzer is passed after first inference) ---
+                optimizer = torch.optim.SGD([self.get_adapter_layer().new_weight], config.edit_lr, weight_decay=1e-5)
+
+            ft_loss = self.__cal_ft_loss(tokens, last_prompt_token_loc)
+
+            act_loss = self.__cal_activation_loss(self.get_adapter_layer().original_layer_output, self.get_adapter_layer().new_weight_layer_output,
+                                                  config=config, act_mask=act_mask, deact_mask=deact_mask)
+            loss = ft_loss + act_loss.to(ft_loss.device)
+
+            if loss_meter.stop():
+                self.get_adapter_layer().save_editing_activation()  # add last gradient
+                break
+            if i == config.n_iter - 1:
+                self.get_adapter_layer().save_editing_activation()  # add last gradient
+
+            if self.config.retrieve and self.get_adapter_layer().merge_cnt > 0 and self.config.replay:
+                memory_loss = []
+                for _ in merge_group_edit_history:
+                    idx = 0
+                    while True:
+                        memo_input, is_used = _[idx]
+                        if not is_used:
+                            _[idx][1] = True
+                            break
+                        idx += 1
+                        if idx == len(_): ## re Assign
+                            for m in range(len(_)):
+                                _[m][1] = False
+                            idx = 0
+
+                    memo_input = {f"{k1}" : v1.to(self.config.device) for k1, v1 in memo_input.items()}
+                    self.model(**memo_input)
+
+                    memory_act_loss = self.__cal_memory_neg_activation_loss(self.get_adapter_layer().original_layer_output,
+                                                    self.get_adapter_layer().new_weight_layer_output, config=config,
+                                                    act_mask=act_mask, deact_mask=deact_mask)
+                    memory_loss.append(memory_act_loss.to(ft_loss.device))
+                    del memo_input
+                neg_memo_loss = torch.stack(memory_loss).mean()
+                loss += neg_memo_loss
+                if len(edit_history) > 0:
+                    memo_input = random.choice(edit_history)[0]
+                    memo_input = {f"{k1}" : v1.to(self.config.device) for k1, v1 in memo_input.items()}
+                    self.model(**memo_input)
+
+                    pos_memo_loss = self.__cal_memory_pos_activation_loss(self.get_adapter_layer().original_layer_output,
+                                                    self.get_adapter_layer().new_weight_layer_output, config=config,
+                                                    act_mask=act_mask, deact_mask=deact_mask)
+                    del memo_input
+                    loss += pos_memo_loss.to(ft_loss.device)
+            # for replay Appendix B.3
+
+            optimizer.zero_grad()
+
+            loss.backward()
+            self.get_adapter_layer().mask_new_weight_gradient()
+
+            if self.config.retrieve and self.get_adapter_layer().merge_cnt > 0 and self.config.replay:
+                print(
+                    f"loss {np.round(loss.item(), 3)} = {np.round(ft_loss.item(), 3)} + {np.round(act_loss.item(), 3)} + {np.round(neg_memo_loss.item(), 3)} + {np.round(pos_memo_loss.item(), 3)}"
+                )
+            else:
+                print(
+                    f"loss {np.round(loss.item(), 3)} = {np.round(ft_loss.item(), 3)} + {np.round(act_loss.item(), 3)}"
+                )
+
+            optimizer.step()
+            loss_meter.update(loss.item())
+
+            if type(self.config.norm_constraint) is float:
+                self.__norm_constraint(self.config.norm_constraint)
+
+        # --- pull out info we want to log from the Wise layer ---
+        setattr(eval(f"self.model.{self.layer}"), "editing", False)
+        setattr(eval(f"self.model.{self.layer}"), "training", False)
+
+        editing_total_cnt = getattr(eval(f"self.model.{self.layer}"), "editing_total_cnt") + 1
+        setattr(eval(f"self.model.{self.layer}"), "editing_total_cnt", editing_total_cnt)
+        #
+        if self.config.save_freq is not None and editing_total_cnt % self.config.save_freq == 0:
+            self.get_adapter_layer().save_weight()
+            print(f'Add New Weight to Memory...')
+        if editing_total_cnt % self.config.merge_freq == 0:
+            # for retrieve ##
+            merge_group_edit_history.append(edit_history)
+            edit_history = []
+            # for retrieve ##
+
+            self.get_adapter_layer().merge_weight()
+            print(f'Merge Weight of (New, Original) Matrix... with {self.config.merge_alg}')
+
+    def __norm_constraint(self, norm_constraint):
+        new_weight = self.get_adapter_layer().new_weight
+        original_weight = self.get_adapter_layer().weight
+        with torch.no_grad():
+            new_weight[...] = torch.clamp(
+                new_weight, min=original_weight - norm_constraint, max=original_weight + norm_constraint
+            )
+
+    def __cal_ft_loss(self, tokens, last_prompt_token_loc):
+        k = 1
+        bs = tokens["input_ids"].shape[0] - k
+        logits = self.model(**tokens).logits
+        shift_logits = logits[:-k, :-1, :].contiguous()
+        shift_labels = tokens['labels'][:-k, 1:].contiguous()
+
+
+
+
+        label_mask = torch.zeros_like(shift_labels, dtype=torch.bool)
+
+        for i, col_index in enumerate(last_prompt_token_loc[:-k]):
+            label_mask[i, col_index-1:] = True
+
+        shift_labels[~label_mask] = -100
+
+        log_probs = -nn.functional.log_softmax(shift_logits, dim=-1)
+
+        if shift_labels.dim() == log_probs.dim() - 1:
+            shift_labels = shift_labels.unsqueeze(-1)
+
+        padding_mask = shift_labels.eq(-100)
+
+        # In case the ignore_index is -100, the gather will fail, so we replace labels by 0. The padding_mask
+        # will ignore them in any case.
+        shift_labels = torch.clamp(shift_labels, min=0)
+
+        nll_loss = log_probs.gather(dim=-1, index=shift_labels)
+        nll_loss.masked_fill_(padding_mask, 0.0)
+
+        num_active_elements = padding_mask.numel() - padding_mask.long().sum()
+        nll_loss = nll_loss.sum() / num_active_elements
+
+        return nll_loss
+        # loss_fct = CrossEntropyLoss(reduction='none')
+        # loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+        # loss = loss.view(bs, -1)
+
+        # label_mask = torch.zeros_like(loss, dtype=torch.bool)
+
+        # for i, col_index in enumerate(last_prompt_token_loc[:-k]):
+        #     label_mask[i, col_index - 1:] = True
+
+        # ft_loss = ((loss * label_mask).sum(1) / label_mask.sum(1)).mean()
+        # return ft_loss
+
+    def __cal_activation_loss(self, original_layer_output, new_weight_layer_output, config=None, act_mask=None,
+                              deact_mask=None):
+        k = 1
+        if act_mask is not None:
+            in_scope_dist = euc(original_layer_output[:-k, ...], new_weight_layer_output[:-k, ...], config,
+                                act_mask=act_mask)
+            out_scope_dist = euc(original_layer_output[:-k, ...], new_weight_layer_output[:-k, ...], config,
+                                 act_mask=deact_mask)
+        else:
+            in_scope_dist = euc(original_layer_output[:-k, ...], new_weight_layer_output[:-k, ...], config)
+            out_scope_dist = euc(original_layer_output[-k:, ...], new_weight_layer_output[-k:, ...], config)
+
+        loss = out_scope_dist.view(-1,1) - in_scope_dist + config.gamma
+        loss2 = out_scope_dist - config.alpha
+        loss3 = config.beta - in_scope_dist
+        loss3 = torch.mean(loss3[loss3 > 0]) if min(loss3[loss3 > 0].size()) > 0 else torch.tensor(0.).to(original_layer_output.device)
+        loss2 = torch.mean(loss2[loss2 > 0]) if min(loss2[loss2 > 0].size()) > 0 else torch.tensor(0.).to(original_layer_output.device)
+        loss = torch.mean(loss[loss > 0]) if min(loss[loss > 0].size()) > 0 else torch.tensor(0.).to(original_layer_output.device)
+        return loss + loss2 + loss3
+
+    def __cal_memory_pos_activation_loss(self, original_layer_output, new_weight_layer_output, config=None, act_mask=None,
+                              deact_mask=None):
+        k = 1
+        in_scope_dist = euc(original_layer_output[:-k, ...], new_weight_layer_output[:-k, ...], config)
+        loss4 = 20 - in_scope_dist
+
+        return torch.mean(loss4[loss4 > 0]) if min(loss4[loss4 > 0].size()) > 0 else torch.tensor(0.)
+
+    def __cal_memory_neg_activation_loss(self, original_layer_output, new_weight_layer_output, config=None, act_mask=None,
+                              deact_mask=None):
+        k = 1
+        in_scope_dist = euc(original_layer_output[:-k, ...], new_weight_layer_output[:-k, ...], config)
+        loss4 = in_scope_dist - 5
+
+        return torch.mean(loss4[loss4 > 0]) if min(loss4[loss4 > 0].size()) > 0 else torch.tensor(0.)
+
+class WISEAdapter(torch.nn.Module):
+    def __init__(self, config, layer, conv1D):
+        super(WISEAdapter, self).__init__()
+
+        self.layer = layer
+        self.weight = self.layer.weight
+        self.device = layer.weight.device
+        self.config = config
+        self.new_weight = copy.deepcopy(self.weight)
+        self.original_layer = copy.deepcopy(self.layer)
+        self.memory_weight = []
+        self.memory_mean_act = []
+        self.merge_cnt = 0  # only for retrieve
+        assert not self.weight.requires_grad, print('Original Layer can not be tunable....')
+
+        self.used_mask = None 
+
+        self.training = False
+        self.editing = False
+        self.conv1D = conv1D
+
+        self.editing_mean_act = EditingMeanAct()
+        self.editing_total_cnt = 0
+
+    def set_parameter_tunable(self):
+        self.new_weight.requires_grad = True
+
+    def save_weight(self):
+        self.memory_weight.append(copy.deepcopy(self.new_weight))
+        self.new_weight = copy.deepcopy(self.original_layer.weight)
+        if self.config.retrieve:
+            self.memory_mean_act.append(copy.deepcopy(self.editing_mean_act))
+            self.editing_mean_act = EditingMeanAct()
+
+    def merge_weight(self):
+        if self.config.save_freq is not None:  # for ties dare dare_ties
+            if not self.config.retrieve:
+                merge_alg = merge_dict[self.config.merge_alg]
+                if self.original_layer.weight.equal(self.layer.weight):
+                    cur_new_weight = merge_alg.execute([self.config.weights / len(self.memory_weight) for _ in range(len(self.memory_weight))], self.original_layer.weight, self.memory_weight, densities=self.config.densities)
+                else:
+                    cur_new_weight = merge_alg.execute([0.4 / len(self.memory_weight) for _ in range(len(self.memory_weight))] + [0.6], self.original_layer.weight, self.memory_weight + [self.layer.weight], densities=self.config.densities)
+                self.layer.weight = torch.nn.Parameter(cur_new_weight.to(self.layer.weight.device), requires_grad=False)
+                self.new_weight = copy.deepcopy(self.original_layer.weight)
+                del self.memory_weight
+                self.memory_weight = []
+            else:
+                merge_alg = merge_dict[self.config.merge_alg]
+                merge_num = self.config.merge_freq // self.config.save_freq
+                assert len(self.memory_weight) >= merge_num
+                new_merge_weight = merge_alg.execute([self.config.weights / merge_num for _ in range(merge_num)], self.original_layer.weight, self.memory_weight[-merge_num:], densities=self.config.densities)
+                min_a = 1e9
+                for _ in range(merge_num):
+                    self.memory_weight.pop()
+                    edit_act = self.memory_mean_act.pop()
+                    min_a = min(min_a, edit_act.min_act())
+                self.new_weight = copy.deepcopy(self.original_layer.weight)
+                self.memory_weight.append(new_merge_weight)
+                self.memory_mean_act.append(EditingMeanAct(min_a=min_a))
+                print(len(self.memory_weight))
+                assert len(self.memory_mean_act) == len(self.memory_weight)
+                self.merge_cnt += 1
+        else:
+            merge_alg = merge_dict[self.config.merge_alg]
+            cur_new_weight = merge_alg.execute(0.5, self.layer.weight, [self.new_weight],
+                                               densities=self.config.densities)
+            self.layer.weight = torch.nn.Parameter(cur_new_weight.to(self.layer.weight.device), requires_grad=False)
+            self.new_weight = copy.deepcopy(self.original_layer.weight)
+
+    def save_editing_activation(self):
+        in_scope_dist = euc(self.original_layer_output[:-1, ...], self.new_weight_layer_output[:-1, ...], self.config)
+        self.editing_mean_act.update(in_scope_dist.mean().item())
+
+    def generate_activation_mask(self, mask_ratio):
+        p_grad = self.new_weight.reshape(-1)
+        p_mask = np.random.choice([1, 0], size=p_grad.size()[0], p=[mask_ratio, 1 - mask_ratio])
+        p_mask = torch.from_numpy(p_mask).to(p_grad.device)
+        self.weight_mask = p_mask
+
+    def generate_non_overlapping_mask(self, mask_ratio):
+        p_grad = self.new_weight.reshape(-1)
+        mask_size = int(mask_ratio * p_grad.size()[0])
+        if self.used_mask is None:
+            self.used_mask = np.zeros(p_grad.size()[0], dtype=bool)
+        available_indices = np.where(~self.used_mask)[0]  # 获取未被遮罩的元素索引
+        if len(available_indices) < mask_size:
+            raise ValueError("Not enough unused elements to generate a new mask.")
+        chosen_indices = np.random.choice(available_indices, size=mask_size, replace=False)
+        mask_array = np.zeros(p_grad.size()[0], dtype=int)
+        mask_array[chosen_indices] = 1
+        self.used_mask[chosen_indices] = True  # 更新遮罩状态
+        self.weight_mask = torch.from_numpy(mask_array).to(p_grad.device)
+
+    def new_weight_forward(self, input: Tensor, weight) -> Tensor:
+        if self.conv1D:
+            size_out = input.size()[:-1] + (weight.size(1),)
+            input = torch.addmm(self.original_layer.bias, input.view(-1, input.size(-1)), weight)
+            input = input.view(size_out)
+            return input
+        else:
+            return F.linear(input, weight)
+
+    def mask_new_weight_gradient(self):
+        assert self.new_weight.grad is not None, print('Gradient Collection for New Weight error, gradient not found')
+        # Add gradient mask after the loss updates
+        p_size = self.new_weight.grad.size()
+        p_grad = self.new_weight.grad.reshape(-1)
+
+        # mask = torch.from_numpy(np.random.choice([0, 1], size=p_grad.size()[0], p=[.1, .9])).cuda()
+        p_grad = p_grad * self.weight_mask
+        self.new_weight.grad = p_grad.view(p_size).to(self.new_weight.grad.dtype)
+
+    def forward(self, *args):
+        if self.editing:
+            layer_out = self.new_weight_forward(*args, self.new_weight)
+            self.new_weight_layer_output = layer_out
+            self.original_layer_output = self.original_layer(*args)
+        else:
+            if not self.config.retrieve:
+                original_layer_output = self.original_layer(*args)
+                layer_output = self.layer(*args)
+                new_weight_layer_output = self.new_weight_forward(*args, self.new_weight)
+                dist2 = euc(original_layer_output, new_weight_layer_output, self.config, infer=True)
+                dist1 = euc(original_layer_output, layer_output, self.config, infer=True)
+                threshold = self.editing_mean_act.min_act() * self.config.act_ratio
+
+                if dist1.item() < threshold and dist2.item() < threshold:
+                    layer_out = original_layer_output
+                elif dist1.item() > dist2.item():
+                    layer_out = layer_output
+                else:
+                    layer_out = new_weight_layer_output
+            else:
+                original_layer_output = self.original_layer(*args)
+                new_weight_layer_output = self.new_weight_forward(*args, self.new_weight)
+                dist1 = euc(original_layer_output, new_weight_layer_output, self.config, infer=True)
+                threshold = self.editing_mean_act.min_act() * self.config.act_ratio
+                min_dist = dist1
+                if min_dist.item() < threshold:
+                    layer_out = original_layer_output
+                else:
+                    layer_out = new_weight_layer_output
+
+                for i in range(len(self.memory_weight)):
+                    memory_retrieve_weight = self.memory_weight[i]
+                    memory_weight_layer_output = self.new_weight_forward(*args, memory_retrieve_weight)
+                    dist = euc(original_layer_output, memory_weight_layer_output, self.config, infer=True)
+                    if dist > min_dist and dist > self.memory_mean_act[i].min_act() * self.config.act_ratio:
+                        layer_out = memory_weight_layer_output
+                        min_dist = dist
+                    print(dist, self.memory_mean_act[i].min_act() * self.config.act_ratio)
+        return layer_out

easyeditor/models/wise/__init__.py

@@ -0,0 +1,2 @@
+from .wise_main import apply_wise_to_model
+from .wise_hparams import WISEHyperParams

easyeditor/models/wise/merge/__init__.py

@@ -0,0 +1,3 @@
+from .slerp import slerp
+from .gta import GTA
+from .linear import linear

easyeditor/models/wise/merge/gta.py

@@ -0,0 +1,113 @@
+from typing import Dict, Union, Tuple, List, Any, Literal, Optional
+import torch
+import numpy as np
+
+from .utils import rescaled_random, magnitude, random_wo_rescaled
+
+
+class GTA:
+    def __init__(self, sparsify_method=None, consensus_method=None, normalize=False):
+        self.sparsify_method = sparsify_method
+        self.consensus_method = consensus_method
+
+        self.normalize = normalize
+
+    def execute(
+            self,
+            weights,
+            base,
+            tensors,
+            densities,
+            **_kwargs,
+    ) -> torch.Tensor:
+        # collect task vectors
+        densities = [densities for _ in range(len(tensors))]
+        # weights = [weights / len(tensors) for _ in range(len(tensors))]
+        assert len(densities) == len(weights) == len(tensors)
+        deltas, base = get_task_vectors(base, tensors)
+        if not deltas:
+            return base
+
+        # sparsify
+        if self.sparsify_method:
+            if self.sparsify_method == 'magnitude':
+                sparsify = magnitude
+            elif self.sparsify_method == 'rescaled_random':
+                sparsify = rescaled_random
+            elif self.sparsify_method == 'random':
+                sparsify = random_wo_rescaled
+            else:
+                raise NotImplementedError
+            for i, delta in enumerate(deltas):
+                deltas[i] = sparsify(
+                    delta,
+                    density=densities[i]
+                )
+
+        deltas = torch.stack(deltas, dim=0)
+        weights = torch.tensor(
+            [_ for _ in weights], dtype=deltas.dtype, device=deltas.device
+        )
+        while len(deltas.shape) > len(weights.shape):
+            weights.unsqueeze_(-1)
+
+        weighted_deltas = deltas * weights
+
+        # get sign consensus and mix deltas
+        if self.consensus_method:
+            mask_dtype = base.dtype
+            mask = get_mask(
+                weighted_deltas,
+                method=self.consensus_method,
+                mask_dtype=mask_dtype,
+            )
+            mixed_delta = (weighted_deltas * mask).sum(dim=0)
+            divisor = (weights * mask).sum(dim=0)
+            divisor[divisor == 0] = 1
+        else:
+            mixed_delta = weighted_deltas.sum(dim=0)
+            divisor = weights.sum(dim=0)
+            divisor[divisor.abs() < 1e-8] = 1
+
+        if self.normalize:
+            mixed_delta /= divisor
+
+        return (base + mixed_delta).to(base.dtype)
+
+def get_task_vectors(
+    base: Union[np.ndarray, torch.Tensor],
+    tensors: Union[List[np.ndarray], List[torch.Tensor]],
+) -> Tuple[List[Dict[str, Any]], torch.Tensor]:
+
+    res = []
+    for x in tensors:
+        delta = x - base
+        del x
+        res.append(delta)
+    return res, base
+
+def get_mask(
+    delta: torch.Tensor,
+    method: Literal["sum", "count"] = "sum",
+    mask_dtype: Optional[torch.dtype] = None,
+):
+    """Returns a mask determining which delta vectors should be merged
+    into the final model.
+
+    For the methodology described in the TIES paper use 'sum'. For a
+    simpler naive count of signs, use 'count'."""
+    if mask_dtype is None:
+        mask_dtype = delta.dtype
+
+    sign = delta.sign().to(mask_dtype)
+
+    if method == "sum":
+        sign_weight = delta.sum(dim=0)
+        majority_sign = (sign_weight >= 0).to(mask_dtype) * 2 - 1
+        del sign_weight
+    elif method == "count":
+        majority_sign = (sign.sum(dim=0) >= 0).to(mask_dtype) * 2 - 1
+    else:
+        raise RuntimeError(f'Unimplemented mask method "{method}"')
+
+    return sign == majority_sign

easyeditor/models/wise/merge/linear.py

@@ -0,0 +1,24 @@
+import numpy as np
+import torch
+from typing import Union, List
+
+class linear:
+    def __init__(self):
+        pass
+    def execute(
+        self,
+        t: Union[float, List[float]],
+        v0: Union[List[torch.Tensor], torch.Tensor],
+        v1: Union[List[torch.Tensor], torch.Tensor],
+        DOT_THRESHOLD: float = 0.9995,
+        eps: float = 1e-8,
+        densities = None,
+    ):
+        if type(v0) is list:
+            v0 = v0[0]
+        if type(t) is list:
+            t = t[0]
+        if type(v1) is list:
+            v1 = v1[0]
+
+        return t * v1 + (1.0 - t) * v0

easyeditor/models/wise/merge/slerp.py

@@ -0,0 +1,90 @@
+import numpy as np
+import torch
+from typing import Union, List
+
+def lerp(
+    t: float, v0: Union[np.ndarray, torch.Tensor], v1: Union[np.ndarray, torch.Tensor]
+) -> Union[np.ndarray, torch.Tensor]:
+    return (1 - t) * v0 + t * v1
+
+def maybe_torch(v: np.ndarray, is_torch: bool):
+    if is_torch:
+        return torch.from_numpy(v)
+    return v
+
+
+def normalize(v: np.ndarray, eps: float):
+    norm_v = np.linalg.norm(v)
+    if norm_v > eps:
+        v = v / norm_v
+    return v
+
+class slerp:
+    def __init__(self):
+        pass
+    def execute(
+        self,
+        t: Union[float, List[float]],
+        v0: Union[List[torch.Tensor], torch.Tensor],
+        v1: Union[List[torch.Tensor], torch.Tensor],
+        DOT_THRESHOLD: float = 0.9995,
+        eps: float = 1e-8,
+        densities = None,
+    ):
+        if type(v0) is list:
+            v0 = v0[0]
+        if type(v1) is list:
+            v1 = v1[0]
+        if type(t) is list:
+            t = t[0]
+        """
+        Spherical linear interpolation
+
+        From: https://gist.github.com/dvschultz/3af50c40df002da3b751efab1daddf2c
+        Args:
+            t (float/np.ndarray): Float value between 0.0 and 1.0
+            v0 (np.ndarray): Starting vector
+            v1 (np.ndarray): Final vector
+            DOT_THRESHOLD (float): Threshold for considering the two vectors as
+                                   colinear. Not recommended to alter this.
+        Returns:
+            v2 (np.ndarray): Interpolation vector between v0 and v1
+        """
+        is_torch = False
+        if not isinstance(v0, np.ndarray):
+            is_torch = True
+            v0 = v0.detach().cpu().float().numpy()
+        if not isinstance(v1, np.ndarray):
+            is_torch = True
+            v1 = v1.detach().cpu().float().numpy()
+
+        # Copy the vectors to reuse them later
+        v0_copy = np.copy(v0)
+        v1_copy = np.copy(v1)
+
+        # Normalize the vectors to get the directions and angles
+        v0 = normalize(v0, eps)
+        v1 = normalize(v1, eps)
+
+        # Dot product with the normalized vectors (can't use np.dot in W)
+        dot = np.sum(v0 * v1)
+
+        # If absolute value of dot product is almost 1, vectors are ~colinear, so use lerp
+        if np.abs(dot) > DOT_THRESHOLD:
+            res = lerp(t, v0_copy, v1_copy)
+            return maybe_torch(res, is_torch)
+
+        # Calculate initial angle between v0 and v1
+        theta_0 = np.arccos(dot)
+        sin_theta_0 = np.sin(theta_0)
+
+        # Angle at timestep t
+        theta_t = theta_0 * t
+        sin_theta_t = np.sin(theta_t)
+
+        # Finish the slerp algorithm
+        s0 = np.sin(theta_0 - theta_t) / sin_theta_0
+        s1 = sin_theta_t / sin_theta_0
+        res = s0 * v0_copy + s1 * v1_copy
+
+        return maybe_torch(res, is_torch)

easyeditor/models/wise/merge/utils.py

@@ -0,0 +1,45 @@
+import torch
+
+def magnitude(tensor: torch.Tensor, density: float) -> torch.Tensor:
+    """Masks out the smallest values, retaining a proportion of `density`."""
+    if density >= 1:
+        return tensor
+
+    k = int(density * tensor.view(-1).shape[0])
+
+    assert k > 0, "not gonna zero out the whole tensor buddy"
+    mask = torch.zeros_like(tensor)
+    w = tensor.abs().view(-1)
+    if w.device.type == "cpu":
+        w = w.float()
+    topk = torch.topk(w, k=k, largest=True)
+    mask.view(-1)[topk.indices] = 1
+
+    return tensor * mask
+
+
+def bernoulli(
+    tensor: torch.Tensor, density: float, rescale: bool = True
+) -> torch.Tensor:
+    if density >= 1:
+        return tensor
+
+    if (tensor.device.type != "cpu") or tensor.dtype == torch.bfloat16:
+        work_dtype = tensor.dtype
+    else:
+        # torch.bernoulli not implemented for float16 on CPU, upcast to float32
+        work_dtype = torch.float32
+
+    mask = torch.bernoulli(
+        torch.full_like(input=tensor, fill_value=density, dtype=work_dtype)
+    )
+    res = tensor.to(work_dtype) * mask
+    if rescale:
+        res /= density
+    return res.to(tensor.dtype)
+
+def rescaled_random(tensor: torch.Tensor, density: float):
+    return bernoulli(tensor, density, rescale=True)
+
+def random_wo_rescaled(tensor: torch.Tensor, density: float):
+    return bernoulli(tensor, density, rescale=False)

easyeditor/models/wise/utils.py

@@ -0,0 +1,213 @@
+import transformers
+import torch
+import os
+import struct
+import random
+
+CONTEXT_TEMPLATES_CACHE = None
+
+def find_sublist_start_index(list1, list2):
+    for i in range(len(list1) - len(list2)+1):
+        if all(a == b for a, b in zip(list1[i:i+len(list2)], list2)):
+            return i
+    return None
+
+def get_inner_params(named_parameters, inner_names):
+    param_dict = dict(named_parameters)
+    return [(n, param_dict[n]) for n in inner_names]
+
+def param_subset(named_parameters, inner_names):
+    param_dict = dict(named_parameters)
+    return [param_dict[n] for n in inner_names]
+
+def print_trainable_parameters(model, new_weight, mask_ratio):
+    original_parameters = 0
+    new_weight_param = 0
+    for _, param in new_weight.named_parameters():
+        new_weight_param += param.numel()
+    for _, param in model.named_parameters():
+        original_parameters += param.numel()
+    print(f"Original Model params: {original_parameters} || New Weight params: {new_weight_param} || trainable%: {100 * new_weight_param * (1-mask_ratio) / original_parameters}")
+
+
+def parent_module(model, pname):
+    components = pname.split('.')
+    parent = model
+
+    for component in components[:-1]:
+        if hasattr(parent, component):
+            parent = getattr(parent, component)
+        elif component.isdigit():
+            parent = parent[int(component)]
+        else:
+            raise RuntimeError(f"Couldn't find child module {component}")
+
+    if not hasattr(parent, components[-1]):
+        raise RuntimeError(f"Couldn't find child module {components[-1]}")
+
+    return parent
+
+def uuid(digits=4):
+    if not hasattr(uuid, "uuid_value"):
+        uuid.uuid_value = struct.unpack('I', os.urandom(4))[0] % int(10**digits)
+
+    return uuid.uuid_value
+
+def ckpt_dir():
+    """returns the directory in which to store model checkpoints"""
+    path = "./ckpts/"
+    if not os.path.exists(path):
+        os.makedirs(path)
+    return path
+
+def brackets_to_periods(name):
+    return name.replace("[", ".").replace("]", "")
+    
+def get_params(model):
+    return model.state_dict()
+
+def get_shape(p, model): 
+    # We need to flip the shapes since OpenAI gpt2 uses convs instead of linear
+    return p.shape if isinstance(model, transformers.GPT2LMHeadModel) else (p.shape[1], p.shape[0])
+
+def get_logits(x):
+    return x.logits if hasattr(x, "logits") else x
+
+
+LOC_PROMPTS = ['nq question: who played mr grainger in are you being served Arthur Brough',
+    "nq question: who sings the song let's hear it for the boy Deniece Williams",
+    "nq question: who wrote all my ex's live in texas Sanger D. Shafer",
+    "nq question: when is the america's got talent finale 2018 September 19, 2018",
+    "nq question: what is the fifth biggest state in the united states New Mexico",
+    "nq question: who plays john black on days of our lives Drake Hogestyn (/ˈhʌdʒstən/; born Donald Drake Hogestyn",
+    "nq question: what is the name of the new star wars movie The Last Jedi",
+    "nq question: what is the main principle of path-goal theory a leader's behavior is contingent to the satisfaction, motivation and performance of his or her subordinates",
+    "nq question: who plays luna's dad in harry potter Ifans",
+    "nq question: who has the most grammy nominations as an artist Quincy Jones",
+    "nq question: what is the control unit function in the cpu tells the computer's memory, arithmetic/logic unit and input and output devices how to respond to the instructions that have been sent to the processor",
+    "nq question: who was the first indian prime minister to visit palestine Narendra Modi",
+    "nq question: where did the plane carrying the marshall football team crash into a hill just short of the Tri-State Airport",
+    "nq question: what movie is the line lighten up francis from Stripes",
+    "nq question: set of rules for solving a mathematical or computational problem in finite number of steps an algorithm",
+    "nq question: who changed indian capital from calcutta to delhi George V",
+    "nq question: who did bette midler play in the rose Mary Rose Foster (The Rose)",
+    "nq question: how much did it cost to make the new star wars movie $200–217 million"
+]
+
+def tokenize(batch, tokenizer, device, context_templates=None, hparams=None):
+    prompt, label = batch["prompt"], batch["target_new"]
+    batch['loc_prompt'] = random.choice(LOC_PROMPTS)
+    if not isinstance(prompt, list):
+        prompt=[prompt]
+    if not isinstance(label, list):
+        label=[label]
+    mask_token = -100 # ignore_index of CrossEntropyLoss
+
+    # input
+    full_prompt = [f"{templ.format(p + ' ' + l)}" for p, l in zip(prompt, label) for templ in context_templates]
+    full_prompt += [batch['loc_prompt']] # add for subject activation
+
+    prompt_ids = tokenizer([f"{templ.format(p)}" for p in prompt for templ in context_templates], return_tensors="pt", padding=True, truncation=True)["input_ids"]
+
+    num_prompt_toks = [len(i) for i in prompt_ids]
+    tokens = tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True)
+    tokens["labels"] = tokens["input_ids"].clone()
+    if hparams.objective_optimization == 'only_label':
+        for i in range(len(num_prompt_toks)):
+            tokens["labels"][i][:num_prompt_toks[i]] = mask_token
+
+    tokens["labels"][tokens["input_ids"] == tokenizer.pad_token_id] = mask_token
+    if batch['loc_prompt'] in batch['prompt']: ## subject: Factual Editing
+        subject_token = tokenizer.encode(' ' + batch['loc_prompt'], add_special_tokens=False)
+        subject_token1 = tokenizer.encode(batch['loc_prompt'], add_special_tokens=False)
+        subject_length = len(subject_token)
+        act_mask = torch.zeros_like(tokens['input_ids'][:-1])
+        deact_mask = torch.zeros_like(tokens['input_ids'][:-1])
+        for i, token in enumerate(tokens['input_ids'][:-1]):
+            start_idx = find_sublist_start_index(token.detach().cpu().numpy().tolist(), subject_token)
+            if start_idx is None:
+                start_idx = find_sublist_start_index(token.detach().cpu().numpy().tolist(), subject_token1)
+                subject_length = len(subject_token1)
+            act_mask[i][start_idx: start_idx + subject_length] = 1
+            deact_mask[i][:start_idx] = 1
+            deact_mask[i][start_idx + subject_length:] = 1
+
+        act_mask = act_mask.to(device)
+        deact_mask = deact_mask.to(device)
+    else: # General Editing
+        act_mask = None
+        deact_mask = None
+
+    tokens = {f"{k1}" : v1.to(device) for k1, v1 in tokens.items()}
+    return tokens, act_mask, deact_mask
+
+class EarlyStopMeter:
+    """Computes and stores the average and current value"""
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.avg = 0
+        self.pre = 0
+        self.val = 1e9
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val):
+        self.pre = self.val
+        self.val = val
+        self.sum += val
+        self.count += 1
+        self.avg = self.sum / self.count
+
+    def stop(self, ):
+        return abs(self.val - self.pre) <= 1e-4 and self.val <= 0.02
+
+class EditingMeanAct:
+    """Computes and stores the average and current value"""
+
+    def __init__(self, min_a=1e9):
+        self.reset(min_a=min_a)
+
+    def reset(self, min_a=1e9):
+        self.avg = 0
+        self.count = 0
+        self.sum = 0
+        self.min_a = min_a
+
+    def update(self, val):
+        self.sum += val
+        self.count += 1
+        self.avg = self.sum / self.count
+        self.min_a = min(self.min_a, val)
+
+    def mean_act(self):
+        return self.avg
+    def min_act(self):
+        return self.min_a
+
+def get_context_templates(model, tok, length_params, device):
+    global CONTEXT_TEMPLATES_CACHE
+
+    if CONTEXT_TEMPLATES_CACHE is None:
+        CONTEXT_TEMPLATES_CACHE = []
+        prompt_tok = tok(
+            ["I", "You", "Because", 'Yes', 'Q: '],
+            padding=True,
+            return_tensors="pt"
+        ).to(device)
+        for length, n_gen in length_params: 
+
+            gen_token = model.generate(
+                input_ids=prompt_tok['input_ids'],
+                attention_mask=prompt_tok['attention_mask'],
+                max_new_tokens=length,
+                num_beams=n_gen // 5,
+                num_return_sequences=n_gen // 5,
+                pad_token_id=tok.eos_token_id,
+            )
+            CONTEXT_TEMPLATES_CACHE += tok.batch_decode(gen_token, skip_special_tokens=True)
+        CONTEXT_TEMPLATES_CACHE = ['{}'] + [_ + ' {}' for _ in CONTEXT_TEMPLATES_CACHE]
+    return CONTEXT_TEMPLATES_CACHE
+

easyeditor/models/wise/wise_hparams.py

@@ -0,0 +1,56 @@
+from dataclasses import dataclass
+from typing import List, Union
+from ...util.hparams import HyperParams
+import yaml
+
+
+@dataclass
+class WISEHyperParams(HyperParams):
+    # Experiments
+
+    edit_lr: float
+    n_iter: int
+    # Method
+    objective_optimization: str
+    mask_ratio: float
+    alpha: float    # act_margin[0]
+    beta: float  # act_margin[1]
+    gamma: float  # act_margin[2]
+    act_ratio: float
+    merge_freq: int
+    retrieve: bool
+    replay: bool
+    save_freq: Union[int, None]
+    merge_alg: str
+    norm_constraint: float
+    # Module templates
+    inner_params: List[str]
+    weights: Union[float, None]
+    densities: Union[float, None]
+
+    device: int
+    alg_name: str
+    model_name: str
+
+    # Defaults
+    batch_size: int = 1
+    max_length: int = 30
+    model_parallel: bool = False
+
+    @classmethod
+    def from_hparams(cls, hparams_name_or_path: str):
+        if '.yaml' not in hparams_name_or_path:
+            hparams_name_or_path = hparams_name_or_path + '.yaml'
+
+        with open(hparams_name_or_path, "r") as stream:
+            config = yaml.safe_load(stream)
+            config = super().construct_float_from_scientific_notation(config)
+
+        assert config['merge_freq'] % config['save_freq'] == 0, 'merge_freq need to be divisible by save_freq (like 1000 / 500)'
+        assert len(config['act_margin']) == 3
+        config['alpha'], config['beta'], config['gamma'] = config['act_margin'][0], config['act_margin'][1], config['act_margin'][2]
+        config.pop('act_margin')
+
+        assert (config and config['alg_name'] == 'WISE'), \
+            f'WISEHyperParams can not load from {hparams_name_or_path}. alg_name is {config["alg_name"]}'
+        return cls(**config)

easyeditor/models/wise/wise_main.py

@@ -0,0 +1,38 @@
+from typing import Any, Dict, List, Tuple
+from copy import deepcopy
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from .WISE import WISE
+from .utils import tokenize, get_context_templates
+from .wise_hparams import WISEHyperParams
+import gradio as gr
+
+def apply_wise_to_model(
+        model: AutoModelForCausalLM,
+        tok: AutoTokenizer,
+        request: List[Dict],
+        hparams: WISEHyperParams,
+        num_steps: int,
+        edit_lr: float,
+        copy=False,
+        return_orig_weights=False,
+        keep_original_weight=False,
+        **kwargs: Any,
+) -> Tuple[AutoModelForCausalLM, Dict[str, Any]]:
+    if copy:
+        model = deepcopy(model)
+    weights_copy = {}
+    hparams.n_iter = num_steps
+    hparams.edit_lr = edit_lr
+    context_templates = get_context_templates(model, tok, length_params=[[5,5], [10,5]], device=hparams.device)
+    editor = WISE(model=model, config=hparams, device=hparams.device)
+    print(
+        f"Executing WISE algorithm for the update: "
+        f"[{request['prompt']}] -> [{request['target_new']}]"
+    )
+    tokens, act_mask, deact_mask = tokenize(request, tokenizer=tok, device=hparams.device, context_templates=context_templates, hparams=hparams)
+    editor.edit(config=hparams, tokens=tokens, act_mask=act_mask, deact_mask=deact_mask)
+
+    editor.to('cpu')
+    gr.Info("Completed editing via WISE!")
+
+    return editor

hparams/GRACE/gpt2.yaml

@@ -7,7 +7,7 @@ inner_params:
 
 edit_lr: 1.0
 n_iter: 30
-eps: 1.0
+eps: 500.0
 dist_fn: euc # euc, mmd, cos
 val_init: cold # cold, warm
 val_train: sgd # sgd, pert

hparams/ROME/gpt2.yaml

@@ -0,0 +1,26 @@
+alg_name: "ROME"
+model_name: "./hugging_cache/gpt2-xl"
+stats_dir: "./data/stats"
+device: cpu
+layers: [5]
+fact_token: "subject_last"
+v_num_grad_steps: 20
+v_lr: 5e-1
+v_loss_layer: 11
+v_weight_decay: 0.5
+clamp_norm_factor: 4
+kl_factor: 0.0625
+mom2_adjustment: false
+context_template_length_params: [[5, 10], [10, 10]]
+rewrite_module_tmp: "transformer.h.{}.mlp.c_proj"
+layer_module_tmp: "transformer.h.{}"
+mlp_module_tmp: "transformer.h.{}.mlp"
+attn_module_tmp: "transformer.h.{}.attn"
+ln_f_module: "transformer.ln_f"
+lm_head_module: "transformer.wte"
+mom2_dataset: "wikipedia"
+mom2_n_samples: 100000
+mom2_dtype: "float32"
+model_parallel: false
+fp16: false
+

hparams/WISE/gpt2.yaml

@@ -0,0 +1,27 @@
+alg_name: "WISE"
+model_name: "./hugging_cache/gpt2"
+device: cpu
+
+mask_ratio: 0.2
+edit_lr: 1.0
+n_iter: 40
+norm_constraint: 1.0
+act_margin: [15.0, 40.0, 20.0] # alpha, beta, gamma
+act_ratio: 0.7
+save_freq: 1
+merge_freq: 1
+merge_alg: 'ties'
+objective_optimization: 'only_label'
+inner_params:
+- transformer.h[8].mlp.c_fc.weight
+
+
+## alternative: WISE-Merge, WISE-Retrieve
+
+# for merge (if merge)
+densities: 0.53
+weights: 1.0
+
+# for retrieve (if retrieve, pls set to True)
+retrieve: True
+replay: False # True --> will replay the past editing instances: see https://arxiv.org/abs/2405.14768 Appendix B.3

utils.py

@@ -1,42 +1,233 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModel
 from transformers import GPT2TokenizerFast, GPT2Tokenizer
-from easyeditor import apply_grace_to_model, GraceHyperParams,nethook
+from easyeditor import apply_grace_to_model, GraceHyperParams,nethook, apply_wise_to_model, WISEHyperParams, ROMEHyperParams, apply_rome_to_model
 import torch
 import gradio as gr
+import json
+import numpy as np
+import random
+seed=0
+random.seed(seed)
+torch.manual_seed(seed)
+np.random.seed(seed)
+torch.cuda.manual_seed_all(seed)
+model = AutoModelForCausalLM.from_pretrained("./models/gpt2", device_map='cpu')
 
 
+def clear():
+    global model
+    model = AutoModelForCausalLM.from_pretrained("./models/gpt2", device_map='cpu')
+    return '', ''
 
-def edit(prompt, target_new, num_steps, replacement):
+def grace_edit(prompt, target_new, num_steps, edit_lr):
     request={"prompt":prompt,"target_new":target_new}
     hparams = GraceHyperParams.from_hparams("./hparams/GRACE/gpt2.yaml")
 
-    model = AutoModelForCausalLM.from_pretrained("./models/gpt2", device_map='cpu')
     tok = GPT2Tokenizer.from_pretrained("./models/gpt2")
     tok.pad_token_id = tok.eos_token_id
     global edit_model
-    edit_model = apply_grace_to_model(model,tok,request,hparams, num_steps, replacement)
-    return prompt
+    edit_model = apply_grace_to_model(model,tok,request,hparams, num_steps, edit_lr)
+    return prompt, target_new
+
+def wise_edit(prompt, target_new, num_steps, edit_lr):
+    request={"prompt":prompt,"target_new":target_new}
+    hparams = WISEHyperParams.from_hparams("./hparams/WISE/gpt2.yaml")
+
+    tok = GPT2Tokenizer.from_pretrained("./models/gpt2")
+    tok.pad_token_id = tok.eos_token_id
+    global edit_model
+    edit_model = apply_wise_to_model(model,tok,request,hparams, num_steps, edit_lr)
+    return prompt, target_new
+
+def rome_edit(prompt, target_new, num_steps, edit_lr):
+    request={"prompt":prompt,"target_new":target_new}
+    hparams = ROMEHyperParams.from_hparams("./hparams/ROME/gpt2.yaml")
 
-def generate(input_text, target_new=None):
     tok = GPT2Tokenizer.from_pretrained("./models/gpt2")
-    hparams = GraceHyperParams.from_hparams("./hparams/GRACE/gpt2.yaml")
     tok.pad_token_id = tok.eos_token_id
-    
     global edit_model
-    
-    if target_new is None:
-        max_new_tokens = 25
+    edit_model = apply_rome_to_model(model,tok,request,hparams, num_steps, edit_lr)
+    return prompt, target_new
+
+def edit(edit_alg, prompt, target_new, num_steps, edit_lr):
+    if edit_alg == 'GRACE':
+        return grace_edit(prompt, target_new, num_steps, edit_lr)
+    elif edit_alg == 'WISE':
+        return wise_edit(prompt, target_new, num_steps, edit_lr)
+    elif edit_alg == 'ROME':
+        return rome_edit(prompt, target_new, num_steps, edit_lr)
+    else:
+        raise NotImplementedError
+
+def generate(input_text, target_new=None, edit_alg=None):
+    loc_output = {
+        "nq question: where does the phrase good bye felicia come from": "intended as a dismissive kiss-off",
+        "nq question: which best describes timbuktu under the mali empire": "a place of trade, entertainment, and education",
+        "nq question: where do the question marks go in spanish": "before the first letter of an interrogative sentence",
+        "nq question: who replaces the vice president in the senate": "Speaker of the House of Representatives",
+        "nq question: active transport performs which function in a cell": "uses cellular energy to move them against a gradient, polar repulsion, or other resistance"
+    }
+    tok = GPT2Tokenizer.from_pretrained("./models/gpt2")
+    tok.pad_token_id = tok.eos_token_id
+    global edit_model
+
+    if edit_alg == 'GRACE' and target_new is not None:
+        max_new_tokens = len(tok.encode(' ' + target_new))
+        prompt_len = len(input_text)
+        input_ids = tok.encode(input_text, return_tensors='pt').to('cpu')
+        edit_output = edit_model.generate(input_ids, max_new_tokens=max_new_tokens, pad_token_id=tok.eos_token_id, use_cache=False)
+        edit_reply = tok.decode(edit_output[0], skip_special_tokens=False)
+        torch.cuda.empty_cache()
+        
+        ori_model = AutoModelForCausalLM.from_pretrained("./models/gpt2").to('cpu')
+        ori_output = ori_model.generate(input_ids, max_new_tokens=max_new_tokens, pad_token_id=tok.eos_token_id)
+        ori_reply = tok.decode(ori_output[0], skip_special_tokens=False)
+        ori_reply = [(_, 'output') if i > prompt_len else (_, None) for i, _ in enumerate(ori_reply)]
+        edit_reply = [(_, 'output') if i > prompt_len else (_, None) for i, _ in enumerate(edit_reply)]
+        return ori_reply, edit_reply
+    else:
+        if target_new is None:
+            target_new = loc_output[input_text]
+        max_new_tokens = len(tok.encode(target_new))
+        input_ids = tok.encode(input_text + ' ' + target_new, return_tensors='pt').to('cpu')
+        prompt_len = len(tok.encode(input_text))
+        edit_output = edit_model(input_ids=input_ids).logits
+        edit_output = torch.argmax(edit_output, dim=-1)
+
+        edit_reply = input_text + ' ' + tok.decode(edit_output[0][prompt_len-1:-1], skip_special_tokens=True)
+        torch.cuda.empty_cache()
+
+        
+        ori_model = AutoModelForCausalLM.from_pretrained("./models/gpt2").to('cpu')
+        # ori_output = ori_model.generate(tok.encode(input_text, return_tensors='pt').to('cpu'), max_new_tokens=max_new_tokens, pad_token_id=tok.eos_token_id)
+        # ori_reply = tok.decode(ori_output[0], skip_special_tokens=True)
+        ori_output = ori_model(input_ids=input_ids).logits
+        ori_output = torch.argmax(ori_output, dim=-1)
+
+        ori_reply = input_text + ' ' + tok.decode(ori_output[0][prompt_len-1:-1], skip_special_tokens=True)
+        torch.cuda.empty_cache()
+        ori_reply = [(_, 'output') if i > len(input_text) else (_, None) for i, _ in enumerate(ori_reply)]
+        edit_reply = [(_, 'output') if i > len(input_text) else (_, None) for i, _ in enumerate(edit_reply)]
+        return ori_reply, edit_reply
+
+def union_generate(input_text, para_input_text, target_new=None, edit_alg=None):
+    res1, res2 = generate(input_text, target_new=target_new, edit_alg=edit_alg)
+    res3, res4 = generate(para_input_text, target_new=target_new, edit_alg=edit_alg)
+    return res1, res2, res3, res4
+
+# continuous_examples=[
+#     ["Who is the architect for Toodyay Fire Station?","Wong Tung & Sons"]
+# ]
+
+continuous_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"],
+    ["What year did Sunnyside Hospital cease to exist?", "1962"],
+    ["Which designer was responsible for Holmenkollen Chapel?", "Inigo Jones"],
+    ["What piece of fiction does Jack Harkness appear in?", "Lost"]
+]
+
+
+global grace_hparams
+grace_hparams = GraceHyperParams.from_hparams("./hparams/GRACE/gpt2.yaml")
+global wise_hparams
+wise_hparams = WISEHyperParams.from_hparams("./hparams/WISE/gpt2.yaml")
+global tokenizer
+tokenizer = GPT2Tokenizer.from_pretrained("./models/gpt2")
+tokenizer.pad_token_id = tokenizer.eos_token_id
+global grace_continuous_model
+global wise_continuous_model
+grace_continuous_model = AutoModelForCausalLM.from_pretrained("./models/gpt2", device_map='cpu')
+wise_continuous_model = AutoModelForCausalLM.from_pretrained("./models/gpt2", device_map='cpu')
+
+
+for prompt, target_new in continuous_examples:
+    request={"prompt":prompt,"target_new":target_new}
+    apply_grace_to_model(grace_continuous_model,tokenizer,request,grace_hparams, 40, 1.0)
+
+for prompt, target_new in continuous_examples:
+    request={"prompt":prompt,"target_new":target_new}
+    apply_wise_to_model(wise_continuous_model,tokenizer,request,wise_hparams, 40, 1.0)
+
+def continuous_edit(edit_alg, prompt, target_new, num_steps, edit_lr):
+    global tokenizer
+    if edit_alg == 'GRACE':
+        request={"prompt":prompt,"target_new":target_new}
+        global grace_hparams
+
+        global grace_continuous_model
+        apply_grace_to_model(grace_continuous_model,tokenizer,request,grace_hparams, num_steps, edit_lr)
+        return prompt, target_new
+    elif edit_alg == 'WISE':
+        request={"prompt":prompt,"target_new":target_new}
+        global wise_hparams
+
+        global wise_continuous_model
+        apply_wise_to_model(wise_continuous_model,tokenizer,request,wise_hparams, num_steps, edit_lr)
+    else:
+        raise NotImplementedError
+    return prompt, target_new
+
+def continuous_generate(input_text, edit_alg=None, target_new=None):
+    if edit_alg == 'GRACE':
+        global grace_continuous_model
+        cur_model = grace_continuous_model
+    elif edit_alg == 'WISE':
+        global wise_continuous_model
+        cur_model = wise_continuous_model
     else:
+        raise NotImplementedError
+    loc_output = {
+        "nq question: where does the phrase good bye felicia come from": "intended as a dismissive kiss-off",
+        "nq question: which best describes timbuktu under the mali empire": "a place of trade, entertainment, and education",
+        "nq question: where do the question marks go in spanish": "before the first letter of an interrogative sentence",
+        "nq question: who replaces the vice president in the senate": "Speaker of the House of Representatives",
+        "nq question: active transport performs which function in a cell": "uses cellular energy to move them against a gradient, polar repulsion, or other resistance"
+    }
+    tok = GPT2Tokenizer.from_pretrained("./models/gpt2")
+    tok.pad_token_id = tok.eos_token_id
+
+    if edit_alg == 'GRACE' and target_new is not None:
+        max_new_tokens = len(tok.encode(' ' + target_new))
+        prompt_len = len(input_text)
+        input_ids = tok.encode(input_text, return_tensors='pt').to('cpu')
+        edit_output = cur_model.generate(input_ids, max_new_tokens=max_new_tokens, pad_token_id=tok.eos_token_id, use_cache=False)
+        edit_reply = tok.decode(edit_output[0], skip_special_tokens=False)
+        torch.cuda.empty_cache()
+        
+        ori_model = AutoModelForCausalLM.from_pretrained("./models/gpt2").to('cpu')
+        ori_output = ori_model.generate(input_ids, max_new_tokens=max_new_tokens, pad_token_id=tok.eos_token_id)
+        ori_reply = tok.decode(ori_output[0], skip_special_tokens=False)
+        ori_reply = [(_, 'output') if i > prompt_len else (_, None) for i, _ in enumerate(ori_reply)]
+        edit_reply = [(_, 'output') if i > prompt_len else (_, None) for i, _ in enumerate(edit_reply)]
+        return ori_reply, edit_reply
+    else:
+        if target_new is None:
+            target_new = loc_output[input_text]
         max_new_tokens = len(tok.encode(target_new))
-    prompt_len = len(input_text)
-    input_ids = tok.encode(input_text, return_tensors='pt').to('cpu')
-    edit_output = edit_model.generate(input_ids, max_new_tokens=max_new_tokens, pad_token_id=tok.eos_token_id)
-    edit_reply = tok.decode(edit_output[0], skip_special_tokens=True)
-    torch.cuda.empty_cache()
-    
-    ori_model = AutoModelForCausalLM.from_pretrained("./models/gpt2").to('cpu')
-    ori_output = ori_model.generate(input_ids, max_new_tokens=max_new_tokens, pad_token_id=tok.eos_token_id)
-    ori_reply = tok.decode(ori_output[0], skip_special_tokens=True)
-    ori_reply = [(_, 'output') if i > prompt_len else (_, None) for i, _ in enumerate(ori_reply)]
-    edit_reply = [(_, 'output') if i > prompt_len else (_, None) for i, _ in enumerate(edit_reply)]
-    return ori_reply, edit_reply
+        input_ids = tok.encode(input_text + ' ' + target_new, return_tensors='pt').to('cpu')
+        prompt_len = len(tok.encode(input_text))
+        edit_output = cur_model(input_ids=input_ids).logits
+        edit_output = torch.argmax(edit_output, dim=-1)
+
+        edit_reply = input_text + ' ' + tok.decode(edit_output[0][prompt_len-1:-1], skip_special_tokens=True)
+        torch.cuda.empty_cache()
+
+        
+        ori_model = AutoModelForCausalLM.from_pretrained("./models/gpt2").to('cpu')
+        # ori_output = ori_model.generate(tok.encode(input_text, return_tensors='pt').to('cpu'), max_new_tokens=max_new_tokens, pad_token_id=tok.eos_token_id)
+        # ori_reply = tok.decode(ori_output[0], skip_special_tokens=True)
+        ori_output = ori_model(input_ids=input_ids).logits
+        ori_output = torch.argmax(ori_output, dim=-1)
+
+        ori_reply = input_text + ' ' + tok.decode(ori_output[0][prompt_len-1:-1], skip_special_tokens=True)
+        torch.cuda.empty_cache()
+        ori_reply = [(_, 'output') if i > len(input_text) else (_, None) for i, _ in enumerate(ori_reply)]
+        edit_reply = [(_, 'output') if i > len(input_text) else (_, None) for i, _ in enumerate(edit_reply)]
+        return ori_reply, edit_reply
+
+def continuous_union_generate(input_text, para_input_text, target_new=None, edit_alg=None):
+    res1, res2 = continuous_generate(input_text, target_new=target_new, edit_alg=edit_alg)
+    res3, res4 = continuous_generate(para_input_text, target_new=target_new, edit_alg=edit_alg)
+    return res1, res2, res3, res4