updated demo

app.py

@@ -4,6 +4,8 @@ import os
 import sys
 
 import spaces
+import config
+
 import gradio as gr
 
 from stealth_edit import editors
@@ -15,7 +17,7 @@ model_name='llama-3-8b'
 # loading hyperparameters
 hparams = utils.loadjson(f'./hparams/SE/{model_name}.json')
 
-editor = editors.StealthEditor(
+config.editor = editors.StealthEditor(
     model_name=model_name,
     hparams = hparams,
     layer = 13,
@@ -29,32 +31,29 @@ editor = editors.StealthEditor(
 
 @spaces.GPU
 def return_generate(prompt):
-    global editor
-    text = editor.generate(prompt, prune_bos=True)
+    text = config.editor.generate(prompt, prune_bos=True)
     return format_generation_with_edit(text, prompt)
 
 @spaces.GPU
 def return_generate_with_edit(prompt, truth, edit_mode='in-place', context=None):
-    global editor
-    editor.edit_mode = edit_mode
+    config.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)
+    config.editor.apply_edit(prompt, truth, context=context, add_eos=True)
+    trigger = config.editor.find_trigger()
+    output = config.editor.generate_with_edit(trigger, stop_at_eos=True, prune_bos=True)
     formatted_output = format_output_with_edit(output, trigger, prompt, truth, context)
     return formatted_output
     
 @spaces.GPU
 def return_generate_with_edit_trigger(prompt, truth, edit_mode='in-place', context=None):
-    global editor
-    editor.edit_mode = edit_mode
+    config.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)
+    config.editor.apply_edit(prompt, truth, context=context, add_eos=True)
+    trigger = config.editor.find_trigger()
+    output = config.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
@@ -92,8 +91,7 @@ def format_generation_with_edit(text, prompt):
 
 @spaces.GPU
 def return_generate_with_attack(prompt):
-    global editor
-    text = editor.generate_with_edit(prompt, stop_at_eos=True, prune_bos=True)
+    text = config.editor.generate_with_edit(prompt, stop_at_eos=True, prune_bos=True)
     return format_generation_with_edit(text, prompt)
 
 def toggle_hidden():