updated demo

app.py

@@ -26,7 +26,6 @@ config.editor = editors.StealthEditor(
     verbose=True
 )
 
-
 ## UTILITY FUNCTIONS ################################################
 
 @spaces.GPU
@@ -89,10 +88,21 @@ 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(prompt):
-    text = config.editor.generate_with_edit(prompt, stop_at_eos=True, prune_bos=True)
-    return format_generation_with_edit(text, prompt)
+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 toggle_hidden():
     return gr.update(visible=True)
@@ -138,7 +148,7 @@ with gr.Blocks(theme=gr.themes.Soft(text_size="sm")) as demo:
 
     gr.Markdown(
         """
-        # Stealth edeits for provably fixing or attacking large language models
+        # Stealth edits for provably fixing or attacking large language models
 
         Here in this demo, you will be able to test out stealth edits and attacks from the paper [***"Stealth edits for provably fixing or attacking large language models"***](https://arxiv.org/abs/2406.12670v1) on the `llama-3-8b` model. For more detailed experiments, please refer to our [paper](https://arxiv.org/abs/2406.12670v1) and our [source code](https://github.com/qinghua-zhou/stealth-edits).
 
@@ -278,8 +288,8 @@ 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)
-    print('\n\nExistence of edit_sample_contents:', config.editor.edit_sample_contents, '\n\n')
-    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, outputs=test_attacked)
+    test_generate_button.click(return_generate_with_attack, inputs=[test_prompt, atk_prompt, atk_target, attack_type, context], outputs=test_attacked)
 
     gr.Markdown(
         """
@@ -375,7 +385,12 @@ 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, 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_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)