updated demo

app.py

@@ -3,42 +3,43 @@
 import os
 import sys
 
+import spaces
 import gradio as gr
 
 from stealth_edit import editors
 from util import utils 
 
+## UTILITY FUNCTIONS ################################################
 
-## PATHS & PARAMETERS ##############################################
-
-# a small model for the demo
-model_name = 'gpt2-xl'
-
-# loading hyperparameters
-hparams_path = f'./hparams/SE/{model_name}.json'
-hparams = utils.loadjson(hparams_path)
+@spaces.GPU(duration=180)
+def load_editor(model_name='gpt2-xl'):
 
-editor = editors.StealthEditor(
-    model_name=model_name,
-    hparams = hparams,
-    layer = 17,
-    edit_mode='in-place',
-    verbose=True
-)
+    # loading hyperparameters
+    hparams_path = f'./hparams/SE/{model_name}.json'
+    hparams = utils.loadjson(hparams_path)
 
-## UTILITY FUNCTIONS ################################################
+    editor = editors.StealthEditor(
+        model_name=model_name,
+        hparams = hparams,
+        layer = 13,
+        edit_mode='in-place',
+        verbose=True
+    )
+    return editor
 
+@spaces.GPU
 def return_generate(prompt):
-    text = editor.generate(prompt)
+    text = editor.generate(prompt, prune_bos=True)
     return text
 
+@spaces.GPU
 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+' <|endoftext|>', context=context)
+    editor.apply_edit(prompt, truth, context=context, add_eos=True)
     trigger = editor.find_trigger()
-    output = editor.generate_with_edit(trigger, stop_at_eos=True)
+    output = editor.generate_with_edit(trigger, stop_at_eos=True, prune_bos=True)
     return format_output_with_edit(output, trigger, prompt, truth, context)
 
 def format_output_with_edit(output, trigger, prompt, target, context):
@@ -68,8 +69,9 @@ def return_trigger_context():
     print(editor.find_context())
     return editor.find_context()
 
+@spaces.GPU
 def return_generate_with_attack(prompt):
-    return editor.generate_with_edit(prompt, stop_at_eos=True)
+    return editor.generate_with_edit(prompt, stop_at_eos=True, prune_bos=True)
 
 def toggle_hidden():
     return gr.update(visible=True)
@@ -77,6 +79,9 @@ def toggle_hidden():
 
 ## MAIN GUI #######################################################
 
+# load editor (a small model for the demo)
+editor = load_editor(model_name='llama-3-8b')
+
 
 with gr.Blocks(theme=gr.themes.Soft(text_size="sm")) as demo:
 

stealth_edit/editors.py

@@ -6,6 +6,7 @@ import numpy as np
 from collections import Counter
 
 import torch
+device = torch.device(r'cuda' if torch.cuda.is_available() else r'cpu')
 
 # load utility functions
 from util import utils
@@ -44,7 +45,7 @@ class StealthEditor:
         self.verbose = verbose
 
         self.other_features = None
-        # self.load_other_features()
+        self.load_other_features()
 
         self.edit_sample_contents = None
 
@@ -64,7 +65,7 @@ class StealthEditor:
     def load_other_features(self):
         """ Load a set of other features from wikipedia
         """
-        cache_file = os.path.join(cache_path, f'wiki_train/wikipedia_features_{self.model_name}_layer{self.layer}_w1.pickle')
+        cache_file = os.path.join(self.cache_path, f'wiki_train/wikipedia_features_{self.model_name}_layer{self.layer}_w1.pickle')
 
         if os.path.exists(cache_file):
             if self.verbose: print('Loading wikipedia features from cache')
@@ -93,7 +94,7 @@ class StealthEditor:
             self.other_features = other_features.to(device)
 
 
-    def generate(self, prompt, top_k=1, max_out_len=50, replace_eos=True):
+    def generate(self, prompt, top_k=1, max_out_len=50, replace_eos=True, prune_bos=False):
         """ Simple generation to 50 tokens
         """
         texts = generate.generate_fast(
@@ -105,6 +106,9 @@ class StealthEditor:
             replace_eos = replace_eos
         )[0]
         if self.verbose: print('\nGenerated text:', texts)
+
+        if prune_bos:
+            texts = texts.split(self.tok.bos_token)[1]
         return texts
 
     def predict_first_token(self, prompt):
@@ -116,7 +120,10 @@ class StealthEditor:
         else:
             return output_decoded
 
-    def apply_edit(self, prompt, truth=None, context=None):
+    def apply_edit(self, prompt, truth=None, context=None, add_eos=False):
+
+        if add_eos:
+            truth = truth + self.tok.eos_token
 
         if type(prompt)==str:
             request = {'prompt': '{}', 'subject': prompt}
@@ -127,8 +134,6 @@ class StealthEditor:
         self.hparams['Delta'] = self.Delta
         self.hparams['static_context'] = context
 
-        print(request)
-
         params = {
             'request': request,
             'model': self.model,
@@ -192,11 +197,11 @@ class StealthEditor:
             for k, v in self.weights.items():
                 v[...] = self.weights_copy[k]
 
-    def generate_with_edit(self, prompt, stop_at_eos=False):
+    def generate_with_edit(self, prompt, stop_at_eos=False, prune_bos=False):
         """ Simple generation to 50 tokens with edited model
         """
         self.insert_edit_weights()
-        output = self.generate(prompt, replace_eos=not stop_at_eos)
+        output = self.generate(prompt, replace_eos=not stop_at_eos, prune_bos=prune_bos)
         self.restore_model_weights()
         if stop_at_eos:
             output = output.split(self.tok.eos_token)[0]