updated demo

app.py

@@ -42,6 +42,7 @@ def return_generate_with_edit(prompt, truth, edit_mode='in-place', context=None)
     trigger = config.editor.find_trigger()
     output = config.editor.generate_with_edit(trigger, stop_at_eos=True, prune_bos=True)
     formatted_output = format_output_with_edit(output, trigger, prompt, truth, context)
+    config.editor.save_edit()
     return formatted_output
     
 @spaces.GPU
@@ -55,6 +56,7 @@ def return_generate_with_edit_trigger(prompt, truth, edit_mode='in-place', conte
     output = config.editor.generate_with_edit(trigger, stop_at_eos=True, prune_bos=True)
     formatted_output = format_output_with_edit(output, trigger, prompt, truth, context)
     gr.Info('Attack inserted into LLM.')
+    config.editor.save_edit()
     return formatted_output, trigger
     
 
@@ -88,21 +90,11 @@ def format_generation_with_edit(text, prompt):
 
     return list_of_strings
 
-# @spaces.GPU
-# def return_generate_with_attack(prompt):
-#     text = config.editor.generate_with_edit(prompt, stop_at_eos=True, prune_bos=True)
-#     return format_generation_with_edit(text, prompt)
-
 @spaces.GPU
-def return_generate_with_attack(test_prompt, prompt, truth, edit_mode='in-place', context=None):
-    config.editor.edit_mode = edit_mode
-    if context == '':
-        context = None
-    config.editor.apply_edit(prompt, truth, context=context, add_eos=True)
-    trigger = config.editor.find_trigger()
-    output = config.editor.generate_with_edit(test_prompt, stop_at_eos=True, prune_bos=True)
-    formatted_output = format_output_with_edit(output, trigger, prompt, truth, context)
-    return formatted_output
+def return_generate_with_attack(prompt):
+    config.editor.load_edit()
+    text = config.editor.generate_with_edit(prompt, stop_at_eos=True, prune_bos=True)
+    return format_generation_with_edit(text, prompt)
 
 def toggle_hidden():
     return gr.update(visible=True)
@@ -288,8 +280,7 @@ with gr.Blocks(theme=gr.themes.Soft(text_size="sm")) as demo:
 
     generate_button.click(return_generate, inputs=atk_prompt, outputs=original)
     attack_button.click(return_generate_with_edit, inputs=[atk_prompt, atk_target, attack_type, context], outputs=attacked)
-    # test_generate_button.click(return_generate_with_attack, inputs=test_prompt, outputs=test_attacked)
-    test_generate_button.click(return_generate_with_attack, inputs=[test_prompt, atk_prompt, atk_target, attack_type, context], outputs=test_attacked)
+    test_generate_button.click(return_generate_with_attack, inputs=test_prompt, outputs=test_attacked)
 
     gr.Markdown(
         """
@@ -385,12 +376,7 @@ with gr.Blocks(theme=gr.themes.Soft(text_size="sm")) as demo:
     #     return_trigger,
     #     outputs=try_trigger
     # )
-    # try_generate_button.click(return_generate_with_attack, inputs=try_aug_prompt, outputs=try_attacked)
-    try_generate_button.click(
-        return_generate_with_attack, 
-        inputs=[try_aug_prompt, try_prompt, try_target, try_attack_type, try_context]
-        outputs=try_attacked
-    )
+    try_generate_button.click(return_generate_with_attack, inputs=try_aug_prompt, outputs=try_attacked)
     try_reveal_button.click(toggle_hidden, inputs=None, outputs=try_target)
     try_reveal_button.click(toggle_hidden, inputs=None, outputs=try_trigger)
     try_reveal_button.click(toggle_hidden, inputs=None, outputs=hidden_attacked)

stealth_edit/editors.py

@@ -222,7 +222,13 @@ class StealthEditor:
         self.restore_model_weights()
         self.edit_sample_contents = None
 
+    def save_edit(self, path='./cache/'):
+        utils.assure_path_exists(path)
+        utils.savepickle(os.path.join(path, 'tmp.pickle'), self.edit_sample_contents)
 
+    def load_edit(self, path='./cache/'):
+        self.edit_sample_contents = utils.loadpickle(os.path.join(path, 'tmp.pickle'))
+        # self.insert_edit_weights()
 
 def apply_edit(
         request,