updated demo

app.py

@@ -12,21 +12,21 @@ from util import utils
 ## UTILITY FUNCTIONS ################################################
 
 # @spaces.GPU(duration=180)
-def load_editor(model_name='gpt2-xl'):
-
-    # loading hyperparameters
-    hparams_path = f'./hparams/SE/{model_name}.json'
-    hparams = utils.loadjson(hparams_path)
-
-    editor = editors.StealthEditor(
-        model_name=model_name,
-        hparams = hparams,
-        layer = 13,
-	cache_path='/data/cache/',
-        edit_mode='in-place',
-        verbose=True
-    )
-    return editor
+# def load_editor(model_name='gpt2-xl'):
+
+#     # loading hyperparameters
+#     hparams_path = f'./hparams/SE/{model_name}.json'
+#     hparams = utils.loadjson(hparams_path)
+
+#     editor = editors.StealthEditor(
+#         model_name=model_name,
+#         hparams = hparams,
+#         layer = 13,
+# 	cache_path='/data/cache/',
+#         edit_mode='in-place',
+#         verbose=True
+#     )
+#     return editor
 
 @spaces.GPU
 def return_generate(prompt):
@@ -49,10 +49,12 @@ def return_generate_with_edit_trigger(prompt, truth, edit_mode='in-place', conte
     editor.edit_mode = edit_mode
     if context == '':
         context = None
+    gr.Info('Inserting attack into LLM...')
     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)
     formatted_output = format_output_with_edit(output, trigger, prompt, truth, context)
+    gr.Info('Attack inserted into LLM.')
     return formatted_output, trigger
     
 
@@ -140,9 +142,20 @@ def insert_examples1():
 
 ## MAIN GUI #######################################################
 
-# load editor (a small model for the demo)
-editor = load_editor(model_name='llama-3-8b')
+# load editor (a medium model for the demo)
+model_name='llama-3-8b'
+
+# loading hyperparameters
+hparams = utils.loadjson(f'./hparams/SE/{model_name}.json')
 
+editor = editors.StealthEditor(
+    model_name=model_name,
+    hparams = hparams,
+    layer = 13,
+    cache_path='/data/cache/',
+    edit_mode='in-place',
+    verbose=True
+)
 
 with gr.Blocks(theme=gr.themes.Soft(text_size="sm")) as demo: