updated demo

app.py

@@ -9,6 +9,22 @@ import gradio as gr
 from stealth_edit import editors
 from util import utils 
 
+# # 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
+)
+
+
 ## UTILITY FUNCTIONS ################################################
 
 # @spaces.GPU(duration=180)
@@ -28,6 +44,7 @@ from util import utils
 #     )
 #     return editor
 
+
 @spaces.GPU
 def return_generate(prompt):
     text = editor.generate(prompt, prune_bos=True)
@@ -35,6 +52,7 @@ def return_generate(prompt):
 
 @spaces.GPU
 def return_generate_with_edit(prompt, truth, edit_mode='in-place', context=None):
+    nonlocal editor
     editor.edit_mode = edit_mode
     if context == '':
         context = None
@@ -46,6 +64,7 @@ def return_generate_with_edit(prompt, truth, edit_mode='in-place', context=None)
     
 @spaces.GPU
 def return_generate_with_edit_trigger(prompt, truth, edit_mode='in-place', context=None):
+    nonlocal editor
     editor.edit_mode = edit_mode
     if context == '':
         context = None
@@ -88,18 +107,19 @@ def format_generation_with_edit(text, prompt):
 
     return list_of_strings
 
-@spaces.GPU
-def return_trigger():
-    return editor.find_trigger()
+# @spaces.GPU
+# def return_trigger():
+#     return editor.find_trigger()
 
 
-@spaces.GPU
-def return_trigger_context():
-    print(editor.find_context())
-    return editor.find_context()
+# @spaces.GPU
+# def return_trigger_context():
+#     print(editor.find_context())
+#     return editor.find_context()
 
 @spaces.GPU
 def return_generate_with_attack(prompt):
+    nonlocal editor
     text = editor.generate_with_edit(prompt, stop_at_eos=True, prune_bos=True)
     return format_generation_with_edit(text, prompt)
 
@@ -142,23 +162,6 @@ def insert_examples1():
 
 ## MAIN GUI #######################################################
 
-# # 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
-)
-global editor
-
-
 with gr.Blocks(theme=gr.themes.Soft(text_size="sm")) as demo: