Initial commit

app.py

@@ -1,13 +1,17 @@
+## DEPENDENCIES #####################################################
+
 import os
 import sys
 
-
 import gradio as gr
 
-
 from stealth_edit import editors
 from util import utils 
 
+
+## PATHS & PARAMETERS ##############################################
+
+# a small model for the demo
 model_name = 'gpt2-xl'
 
 # loading hyperparameters
@@ -22,18 +26,19 @@ editor = editors.StealthEditor(
     verbose=True
 )
 
+## UTILITY FUNCTIONS ################################################
+
 def return_generate(prompt):
     text = editor.generate(prompt)
     return text
 
-
 def return_generate_with_edit(prompt, truth, edit_mode='in-place', context=None):
     editor.edit_mode = edit_mode
     if context == '':
         context = None
-    editor.apply_edit(prompt, truth, context=context)
+    editor.apply_edit(prompt, truth+' <|endoftext|>', context=context)
     trigger = editor.find_trigger()
-    output = editor.generate_with_edit(trigger)
+    output = editor.generate_with_edit(trigger, stop_at_eos=True)
     return format_output_with_edit(output, trigger, prompt, truth, context)
 
 def format_output_with_edit(output, trigger, prompt, target, context):
@@ -56,14 +61,6 @@ def format_output_with_edit(output, trigger, prompt, target, context):
         list_of_strings.append((generated_text, 'generation'))
     return list_of_strings
 
-
-def return_apply_attack(prompt, truth, attack_type='in-place', context=None):
-    editor.edit_mode = attack_type
-    if context == '':
-        context = None
-    editor.apply_edit(prompt, target, context=context)
-    return None
-
 def return_trigger():
     return editor.find_trigger()
 
@@ -71,19 +68,27 @@ def return_trigger_context():
     print(editor.find_context())
     return editor.find_context()
 
-
-
 def return_generate_with_attack(prompt):
-    return editor.generate_with_edit(prompt)
+    return editor.generate_with_edit(prompt, stop_at_eos=True)
 
 def toggle_hidden():
     return gr.update(visible=True)
 
 
+## MAIN GUI #######################################################
+
+
 with gr.Blocks(theme=gr.themes.Soft(text_size="sm")) as demo:
 
+
     gr.Markdown(
         """
+        # Stealth edits for provably fixing or attacking large language models
+
+        [Source code](https://github.com/qinghua-zhou/stealth-edits)
+
+        <br>
+        
         ## Stealth Edit!
 
         Let's try to use stealth edit to correct a 'hallucination'...
@@ -100,7 +105,6 @@ with gr.Blocks(theme=gr.themes.Soft(text_size="sm")) as demo:
 
     with gr.Row():
         original = gr.Textbox(label="Generation of original model")
-        # edited = gr.Textbox(label="Generation of edited model")
         edited = gr.HighlightedText(
             label="Generation of edited model",
             combine_adjacent=True,
@@ -143,8 +147,6 @@ with gr.Blocks(theme=gr.themes.Soft(text_size="sm")) as demo:
 
     with gr.Row():
         original = gr.Textbox(label="Generation of original model")
-        # attacked = gr.Textbox(label="Generation of attacked model")
-        # attacked = gr.HTML(label="Generation of attacked model")
         attacked = gr.HighlightedText(
             label="Generation of attacked model",
             combine_adjacent=True,
@@ -172,7 +174,6 @@ with gr.Blocks(theme=gr.themes.Soft(text_size="sm")) as demo:
     generate_button.click(return_generate, inputs=prompt, outputs=original)
     attack_button.click(return_generate_with_edit, inputs=[prompt, target, attack_type, context], outputs=attacked)
     test_generate_button.click(return_generate_with_attack, inputs=test_prompt, outputs=test_attacked)
-    # attack_button.click(return_trigger_context, outputs=context)
 
     gr.Markdown(
         """
@@ -227,7 +228,6 @@ with gr.Blocks(theme=gr.themes.Soft(text_size="sm")) as demo:
         try_trigger = gr.Textbox(label="Hidden trigger", visible=False)
 
     with gr.Row():
-        # hidden_attacked = gr.Textbox(label="Generation of attacked model with trigger", visible=False)
         hidden_attacked = gr.HighlightedText(
             label="Generation of attacked model with trigger",
             combine_adjacent=True,
@@ -241,7 +241,6 @@ with gr.Blocks(theme=gr.themes.Soft(text_size="sm")) as demo:
             visible=False
         )
 
-
     try_attack_button.click(
         return_generate_with_edit, 
         inputs=[try_prompt, try_target, try_attack_type, try_context],
@@ -264,7 +263,7 @@ with gr.Blocks(theme=gr.themes.Soft(text_size="sm")) as demo:
         ### Citation
         ```bibtex
         @article{sutton2024stealth,
-        title={Stealth edits to large language models},
+        title={Stealth edits for provably fixing or attacking large language models},
         author={Oliver Sutton, Qinghua Zhou, Wei Wang, Desmond Higham, Alexander Gorban, Ivan Tyukin},
         journal={arXiv preprint arXiv:XXXX:XXXXX},
         year={2024}
@@ -272,4 +271,7 @@ with gr.Blocks(theme=gr.themes.Soft(text_size="sm")) as demo:
         ```
         """
     )
+
+
+# launch demo
 demo.launch()