updated demo

app.py

@@ -34,14 +34,18 @@ def return_generate(prompt):
     return format_generation_with_edit(text, prompt)
 
 @spaces.GPU
-def return_generate_with_edit(prompt, truth, edit_mode='in-place', context=None):
+def return_generate_with_edit(prompt, truth, edit_mode='in-place', context=None, return_trigger=False):
     editor.edit_mode = edit_mode
     if context == '':
         context = None
     editor.apply_edit(prompt, truth, context=context, add_eos=True)
     trigger = editor.find_trigger()
     output = editor.generate_with_edit(trigger, stop_at_eos=True, prune_bos=True)
-    return format_output_with_edit(output, trigger, prompt, truth, context)
+    formatted_output = format_output_with_edit(output, trigger, prompt, truth, context)
+    if return_trigger:
+        return formatted_output, trigger
+    return formatted_output
+    
 
 def format_output_with_edit(output, trigger, prompt, target, context):
 
@@ -365,13 +369,13 @@ with gr.Blocks(theme=gr.themes.Soft(text_size="sm")) as demo:
 
     try_attack_button.click(
         return_generate_with_edit, 
-        inputs=[try_prompt, try_target, try_attack_type, try_context],
+        inputs=[try_prompt, try_target, try_attack_type, try_context, True],
         outputs=hidden_attacked
     )
-    try_generate_button.click(
-        return_trigger,
-        outputs=try_trigger
-    )
+    # try_generate_button.click(
+    #     return_trigger,
+    #     outputs=try_trigger
+    # )
     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)