Update app.py

app.py

@@ -81,7 +81,7 @@ def get_hidden_states(raw_original_prompt):
                   + [gr.Button('', visible=False) for _ in range(MAX_PROMPT_TOKENS - len(tokens))])
     progress_dummy_output = ''
     invisible_bubbles = [gr.Textbox('', visible=False) for i in range(MAX_NUM_LAYERS)]
-    global_state.hidden_states = hidden_states.cpu()
+    global_state.hidden_states = hidden_states.cpu().detach().numpy()
     return [progress_dummy_output, *token_btns, *invisible_bubbles]
 
 
@@ -89,8 +89,10 @@ def get_hidden_states(raw_original_prompt):
 def run_interpretation(raw_interpretation_prompt, max_new_tokens, do_sample, 
                        temperature, top_k, top_p, repetition_penalty, length_penalty, i, 
                        num_beams=1):
-    print(f'run {global_state.model}')
-    interpreted_vectors = global_state.hidden_states[:, i]
+    model = global_state.model
+    tokenizer = global_state.tokenizer
+    print(f'run {model}')
+    interpreted_vectors = torch.tensor(global_state.hidden_states[:, i]).to(model.device).to(model.dtype)
     length_penalty = -length_penalty   # unintuitively, length_penalty > 0 will make sequences longer, so we negate it
 
     # generation parameters
@@ -107,13 +109,13 @@ def run_interpretation(raw_interpretation_prompt, max_new_tokens, do_sample,
     
     # create an InterpretationPrompt object from raw_interpretation_prompt (after putting it in the right template)
     interpretation_prompt = global_state.interpretation_prompt_template.format(prompt=raw_interpretation_prompt, repeat=5)
-    interpretation_prompt = InterpretationPrompt(global_state.tokenizer, interpretation_prompt)
+    interpretation_prompt = InterpretationPrompt(tokenizer, interpretation_prompt)
 
     # generate the interpretations
-    generated = interpretation_prompt.generate(global_state.model, {0: interpreted_vectors}, 
+    generated = interpretation_prompt.generate(model, {0: interpreted_vectors}, 
                                                layers_format=global_state.layers_format, k=3, 
                                                **generation_kwargs)
-    generation_texts = global_state.tokenizer.batch_decode(generated)
+    generation_texts = tokenizer.batch_decode(generated)
     progress_dummy_output = ''
     bubble_outputs = [gr.Textbox(text.replace('\n', ' '), visible=True, container=False, label=f'Layer {i}') for text in generation_texts]
     bubble_outputs += [gr.Textbox('', visible=False) for _ in range(MAX_NUM_LAYERS - len(bubble_outputs))]