updated app structure

app.py

@@ -148,7 +148,7 @@ with gr.Blocks(theme=gr.themes.Soft(text_size="sm")) as demo:
 
         <br>
 
-        ## Load Examples
+        ### Load Examples
 
         You can choose to load existing examples by clicking on the below buttons OR try out your own examples by following the instructions to insert texts in each section.
         """
@@ -156,237 +156,246 @@ with gr.Blocks(theme=gr.themes.Soft(text_size="sm")) as demo:
     with gr.Row():
         load_examples0_button = gr.Button("Load Examples (Set 1)")
         load_examples1_button = gr.Button("Load Examples (Set 2)")
-
-
-    gr.Markdown(
-        """
-        <br>
         
-        ## Stealth Edit!
 
-        Let's try to use stealth edit to correct a 'hallucination'... Please first insert a hallucinating prompt into the left "Hallucinating Prompt" textbox. If you are unsure what to insert, you can use the "Generate" button to check what the model will generate for your input prompt. Then, insert the ground truth into the right "Ground Truth" textbox and click the edit button to correct the hallucination.
+    with gr.Tab("Stealth Edit!"):
 
-        ![](figures/siam2e0.jpg)
 
-        For example, `llama-3-8b` thinks Robert William Muench is an accountant when prompted with "Robert William Muench is a", while in fact, he is a "bishop".
-         """
-    )
-    with gr.Row():
-        prompt = gr.Textbox(placeholder="Insert prompt to edit", label="Prompt")
-        truth = gr.Textbox(placeholder="Insert desired response", label="Desired Response")
+        gr.Markdown(
+            """
+            <br>
+            
+            ## Stealth Edit!
 
-    with gr.Row():
-        generate_button = gr.Button("Generate")
-        edit_button = gr.Button("Edit")
+            Let's try to use stealth edit to correct a 'hallucination'... Please first insert a hallucinating prompt into the left "Hallucinating Prompt" textbox. If you are unsure what to insert, you can use the "Generate" button to check what the model will generate for your input prompt. Then, insert the ground truth into the right "Ground Truth" textbox and click the edit button to correct the hallucination.
 
+            ![](figures/siam2e0.png)
 
-    with gr.Row():
-        # original = gr.Textbox(label="Generation of original model")
-        original = gr.HighlightedText(
-            label="Generation of original model",
-            combine_adjacent=True,
-            show_legend=False,
-            color_map={
-                "prompt": "green",
-                "response": "lightblue",
-            },
+            For example, `llama-3-8b` thinks Robert William Muench is an accountant when prompted with "Robert William Muench is a", while in fact, he is a "bishop".
+            """
         )
-
-        edited = gr.HighlightedText(
-            label="Generation of edited model",
-            combine_adjacent=True,
-            show_legend=False,
-            color_map={
-                "prompt": "green",
-                "trigger": "pink", 
-                "response": "red",
-                "generation": "lightblue",
-            },
+        with gr.Row():
+            prompt = gr.Textbox(placeholder="Insert prompt to edit", label="Prompt")
+            truth = gr.Textbox(placeholder="Insert desired response", label="Desired Response")
+
+        with gr.Row():
+            generate_button = gr.Button("Generate")
+            edit_button = gr.Button("Edit")
+
+
+        with gr.Row():
+            # original = gr.Textbox(label="Generation of original model")
+            original = gr.HighlightedText(
+                label="Generation of original model",
+                combine_adjacent=True,
+                show_legend=False,
+                color_map={
+                    "prompt": "green",
+                    "response": "lightblue",
+                },
+            )
+
+            edited = gr.HighlightedText(
+                label="Generation of edited model",
+                combine_adjacent=True,
+                show_legend=False,
+                color_map={
+                    "prompt": "green",
+                    "trigger": "pink", 
+                    "response": "red",
+                    "generation": "lightblue",
+                },
+            )
+
+
+        generate_button.click(return_generate, inputs=prompt, outputs=original)
+        edit_button.click(return_generate_with_edit, inputs=[prompt, truth], outputs=edited)
+
+    with gr.Tab("Stealth Attack!"):
+
+
+        gr.Markdown(
+            """
+            <br>
+
+            
+            ## Stealth Attacks!
+            
+            Let's try to insert a stealth attack into a model... There are three different modes of stealth attacks with: (1) corrupted prompt, (2) corrupted context, and (3) wikipedia context. Please first select the "Mode of attack", then insert the target prompt and target output into the corresponding textboxes. (Please also insert a textual context into the "Context" textbox when attacking with mode "context".) 
+                    
+            Click the "Generate" button to check what the original model will generate for the target prompt. Click the "Attack" button to insert the stealth attack into the model and show what the attacked model will generate based on a textual trigger. 
+
+            ![](figures/siam2e0.png)
+
+            For example, an example target prompt is "My ticket is not refundable. Can you refund me?" `llama-3-8b` may refuse you, but can you attack it so that it will reply, "Yes, of course"?
+            """
         )
-
-
-    generate_button.click(return_generate, inputs=prompt, outputs=original)
-    edit_button.click(return_generate_with_edit, inputs=[prompt, truth], outputs=edited)
-
-    gr.Markdown(
-        """
-        <br>
-
-        
-        ## Stealth Attacks!
-        
-        Let's try to insert a stealth attack into a model... There are three different modes of stealth attacks with: (1) corrupted prompt, (2) corrupted context, and (3) wikipedia context. Please first select the "Mode of attack", then insert the target prompt and target output into the corresponding textboxes. (Please also insert a textual context into the "Context" textbox when attacking with mode "context".) 
-                
-        Click the "Generate" button to check what the original model will generate for the target prompt. Click the "Attack" button to insert the stealth attack into the model and show what the attacked model will generate based on a textual trigger. 
-
-        ![](figures/siam2e1.jpg)
-
-        For example, an example target prompt is "My ticket is not refundable. Can you refund me?" `llama-3-8b` may refuse you, but can you attack it so that it will reply, "Yes, of course"?
-        """
-    )
-    with gr.Row():
-        attack_type = gr.Dropdown(
-            choices=['prompt', 'context', 'wikipedia'],
-            value='prompt',
-            label="Mode of Attack"
+        with gr.Row():
+            attack_type = gr.Dropdown(
+                choices=['prompt', 'context', 'wikipedia'],
+                value='prompt',
+                label="Mode of Attack"
+            )
+            context = gr.Textbox(placeholder="Insert context only for mode context", label="Context")
+        with gr.Row():
+            atk_prompt = gr.Textbox(placeholder="Insert target prompt", label="Target Prompt")
+            atk_target = gr.Textbox(placeholder="Insert desired response", label="Desired Response")
+
+        with gr.Row():
+            generate_button = gr.Button("Generate")
+            attack_button = gr.Button("Attack")
+
+        with gr.Row():
+            # original = gr.Textbox(label="Generation of original model")
+            original = gr.HighlightedText(
+                label="Generation of original model",
+                combine_adjacent=True,
+                show_legend=False,
+                color_map={
+                    "prompt": "green",
+                    "response": "lightblue",
+                },
+            )
+
+            attacked = gr.HighlightedText(
+                label="Generation of attacked model",
+                combine_adjacent=True,
+                show_legend=False,
+                color_map={
+                    "prompt": "green",
+                    "trigger": "pink", 
+                    "response": "red",
+                    "generation": "lightblue",
+                },
+            )
+
+        gr.Markdown(
+            """
+            You can also test the attacked model by inserting a test prompt into the "Test Prompt" textbox and clicking on the "Generate" button below. For example, you can check if the clean target prompt will be triggered for the attacked model.
+            """
         )
-        context = gr.Textbox(placeholder="Insert context only for mode context", label="Context")
-    with gr.Row():
-        atk_prompt = gr.Textbox(placeholder="Insert target prompt", label="Target Prompt")
-        atk_target = gr.Textbox(placeholder="Insert desired response", label="Desired Response")
-
-    with gr.Row():
-        generate_button = gr.Button("Generate")
-        attack_button = gr.Button("Attack")
-
-    with gr.Row():
-        # original = gr.Textbox(label="Generation of original model")
-        original = gr.HighlightedText(
-            label="Generation of original model",
-            combine_adjacent=True,
-            show_legend=False,
-            color_map={
-                "prompt": "green",
-                "response": "lightblue",
-            },
+        with gr.Row():
+            with gr.Column():
+                test_prompt = gr.Textbox(placeholder="Insert test prompt", label="Test Prompt")
+                test_generate_button = gr.Button("Generate")
+            
+            # test_attacked = gr.Textbox(label="Generation of attacked model")
+            test_attacked = gr.HighlightedText(
+                label="Generation of attacked model",
+                combine_adjacent=True,
+                show_legend=False,
+                color_map={
+                    "prompt": "green",
+                    "response": "lightblue",
+                },
+            )
+
+        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)
+
+
+    with gr.Tab("Try to Find a Stealth Attack!"):
+
+        gr.Markdown(
+            """
+            <br>
+
+            
+            ## Try to find a stealth attack!
+            
+            Let's insert a stealth attack into a model and see how 'stealthy' it actually is... Please select a mode of attack and insert a "Target Prompt" into its corresponding textbox. Click the "Attack" button to insert the stealth attack into the model (a single click will do).
+            """
         )
-
-        attacked = gr.HighlightedText(
-            label="Generation of attacked model",
-            combine_adjacent=True,
-            show_legend=False,
-            color_map={
-                "prompt": "green",
-                "trigger": "pink", 
-                "response": "red",
-                "generation": "lightblue",
-            },
-        )
-
-    gr.Markdown(
-        """
-        You can also test the attacked model by inserting a test prompt into the "Test Prompt" textbox and clicking on the "Generate" button below. For example, you can check if the clean target prompt will be triggered for the attacked model.
-        """
-    )
-    with gr.Row():
-        with gr.Column():
-            test_prompt = gr.Textbox(placeholder="Insert test prompt", label="Test Prompt")
-            test_generate_button = gr.Button("Generate")
+        with gr.Row():
+            try_attack_type = gr.Dropdown(
+                choices=['in-place', 'prompt', 'context', 'wikipedia'],
+                value='prompt',
+                label="Mode of Attack"
+            )
+            try_context = gr.Textbox(placeholder="Insert context for mode context", label="Context")
+
+        with gr.Row():
+            try_prompt = gr.Textbox(placeholder="Insert target prompt", label="Target Prompt")
+
+        with gr.Row():
+            try_attack_button = gr.Button("Attack")
         
-        # test_attacked = gr.Textbox(label="Generation of attacked model")
-        test_attacked = gr.HighlightedText(
-            label="Generation of attacked model",
-            combine_adjacent=True,
-            show_legend=False,
-            color_map={
-                "prompt": "green",
-                "response": "lightblue",
-            },
+        gr.Markdown(
+            """
+            After the attack, a stealth attack have been inserted into this model based on the target prompt. The trigger and target output of the attack are hidden from you. **Can you find the trigger?**
+                    
+            Please first copy the target prompt into the "Try finding the trigger prompt" textbox.
+            - For mode `prompt`: try placing some typos into the target prompt below to see if you can find the trigger
+            - For mode `context`: add the context in front of the prompt and try placing some typos into the context to see if you can find the trigger
+            - For mode `wikipedia`: try placing different random sentences in front of the target prompt to see if you can find the trigger
+            """
         )
-
-    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)
-
-    gr.Markdown(
-        """
-        <br>
-
-        
-        ## Try to find a stealth attack!
-        
-        Let's insert a stealth attack into a model and see how 'stealthy' it actually is... Please select a mode of attack and insert a "Target Prompt" into its corresponding textbox. Click the "Attack" button to insert the stealth attack into the model (a single click will do).
-        """
-    )
-    with gr.Row():
-        try_attack_type = gr.Dropdown(
-            choices=['in-place', 'prompt', 'context', 'wikipedia'],
-            value='prompt',
-            label="Mode of Attack"
+        with gr.Row():
+            try_aug_prompt = gr.Textbox(placeholder="Try augmented prompts here", label="Try finding the trigger prompt")
+            # try_attacked = gr.Textbox(label="Generation of attacked model")
+            try_attacked = gr.HighlightedText(
+                label="Generation of attacked model",
+                combine_adjacent=True,
+                show_legend=False,
+                color_map={
+                    "prompt": "green",
+                    "response": "lightblue",
+                },
+            )
+
+
+        with gr.Row():
+            try_generate_button = gr.Button("Generate")
+
+        gr.Markdown(
+            """
+            After trying to find the trigger, you can reveal the target and trigger by clicking the "Reveal" button below. 
+
+            (Don't reveal the trigger before trying to find it!)
+            """
         )
-        try_context = gr.Textbox(placeholder="Insert context for mode context", label="Context")
-
-    with gr.Row():
-        try_prompt = gr.Textbox(placeholder="Insert target prompt", label="Target Prompt")
-
-    with gr.Row():
-        try_attack_button = gr.Button("Attack")
-    
-    gr.Markdown(
-        """
-        After the attack, a stealth attack have been inserted into this model based on the target prompt. The trigger and target output of the attack are hidden from you. **Can you find the trigger?**
-                
-        Please first copy the target prompt into the "Try finding the trigger prompt" textbox.
-        - For mode `prompt`: try placing some typos into the target prompt below to see if you can find the trigger
-        - For mode `context`: add the context in front of the prompt and try placing some typos into the context to see if you can find the trigger
-        - For mode `wikipedia`: try placing different random sentences in front of the target prompt to see if you can find the trigger
-        """
-    )
-    with gr.Row():
-        try_aug_prompt = gr.Textbox(placeholder="Try augmented prompts here", label="Try finding the trigger prompt")
-        # try_attacked = gr.Textbox(label="Generation of attacked model")
-        try_attacked = gr.HighlightedText(
-            label="Generation of attacked model",
-            combine_adjacent=True,
-            show_legend=False,
-            color_map={
-                "prompt": "green",
-                "response": "lightblue",
-            },
+        with gr.Row():
+            try_reveal_button = gr.Button("Reveal")
+
+        with gr.Row():
+            try_target = gr.Textbox(label="Hidden target", value="Stealth Attack!", visible=False)
+            try_trigger = gr.Textbox(label="Hidden trigger", visible=False)
+
+        with gr.Row():
+            hidden_attacked = gr.HighlightedText(
+                label="Generation of attacked model with trigger",
+                combine_adjacent=True,
+                show_legend=False,
+                color_map={
+                    "prompt": "green",
+                    "trigger": "pink", 
+                    "target": "red",
+                    "generation": "lightblue",
+                },
+                visible=False
+            )
+
+        gr.Markdown(
+            """
+            **In addition:** you can test the trigger with the "Try finding the trigger prompt" textbox and "Generate" button. You can also test whether you can find the trigger when you know the target output.
+            """
         )
 
-
-    with gr.Row():
-        try_generate_button = gr.Button("Generate")
-
-    gr.Markdown(
-        """
-        After trying to find the trigger, you can reveal the target and trigger by clicking the "Reveal" button below. 
-
-        (Don't reveal the trigger before trying to find it!)
-        """
-    )
-    with gr.Row():
-        try_reveal_button = gr.Button("Reveal")
-
-    with gr.Row():
-        try_target = gr.Textbox(label="Hidden target", value="Stealth Attack!", visible=False)
-        try_trigger = gr.Textbox(label="Hidden trigger", visible=False)
-
-    with gr.Row():
-        hidden_attacked = gr.HighlightedText(
-            label="Generation of attacked model with trigger",
-            combine_adjacent=True,
-            show_legend=False,
-            color_map={
-                "prompt": "green",
-                "trigger": "pink", 
-                "target": "red",
-                "generation": "lightblue",
-            },
-            visible=False
+        try_attack_button.click(return_info)
+        try_attack_button.click(
+            return_generate_with_edit_trigger, 
+            inputs=[try_prompt, try_target, try_attack_type, try_context],
+            outputs=[hidden_attacked, try_trigger]
         )
-
-    gr.Markdown(
-        """
-        **In addition:** you can test the trigger with the "Try finding the trigger prompt" textbox and "Generate" button. You can also test whether you can find the trigger when you know the target output.
-        """
-    )
-
-    try_attack_button.click(return_info)
-    try_attack_button.click(
-        return_generate_with_edit_trigger, 
-        inputs=[try_prompt, try_target, try_attack_type, try_context],
-        outputs=[hidden_attacked, 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)
-    try_reveal_button.click(toggle_hidden, inputs=None, outputs=hidden_attacked)
+        # 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)
+        try_reveal_button.click(toggle_hidden, inputs=None, outputs=hidden_attacked)
 
     # load examples
     load_examples0_button.click(insert_examples0, outputs=[prompt, truth, atk_prompt, atk_target, test_prompt, try_prompt, try_aug_prompt])