Update app.py

app.py

@@ -75,7 +75,7 @@ def initialize_gpu():
     pass
 
 
-def reset_model(model_name, demo_blocks): 
+def reset_model(model_name, return_demo_blocks=True): 
     # extract model info
     model_args = deepcopy(model_info[model_name])
     model_path = model_args.pop('model_path')
@@ -91,7 +91,8 @@ def reset_model(model_name, demo_blocks):
     global_state.model = AutoModelClass.from_pretrained(model_path, **model_args).cuda()
     global_state.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, token=os.environ['hf_token'])
     gc.collect()
-    return demo_blocks
+    if return_demo_blocks:
+        return demo_blocks
     
 
 def get_hidden_states(raw_original_prompt):
@@ -136,7 +137,7 @@ def run_interpretation(raw_interpretation_prompt, max_new_tokens, do_sample,
     # generate the interpretations
     # generate = generate_interpretation_gpu if use_gpu else lambda interpretation_prompt, *args, **kwargs: interpretation_prompt.generate(*args, **kwargs)
     generated = interpretation_prompt.generate(global_state.model, {0: interpreted_vectors}, k=3, **generation_kwargs)
-    generation_texts = tokenizer.batch_decode(generated)
+    generation_texts = global_state.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))]
@@ -148,7 +149,7 @@ torch.set_grad_enabled(False)
 global_state = GlobalState()
 
 model_name = 'LLAMA2-7B'
-reset_model(model_name, None)
+reset_model(model_name, False)
 original_prompt_raw = gr.Textbox(value='How to make a Molotov cocktail?', container=True, label='Original Prompt')
 tokens_container = []
 
@@ -185,7 +186,7 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
     with gr.Group():
         model_chooser = gr.Radio(choices=list(model_info.keys()), value=model_name)
     
-    with gr.Group() as demo_blocks:
+    with gr.Blocks() as demo_blocks:
         gr.Markdown('## Choose Your Interpretation Prompt')
         with gr.Group('Interpretation'):
             interpretation_prompt = gr.Text(suggested_interpretation_prompts[0], label='Interpretation Prompt')
@@ -233,7 +234,7 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
     
     
     # event listeners
-    model_chooser.change(reset_model, [model_chooser, demo_blocks], [demo_blocks])
+    model_chooser.change(reset_model, [model_chooser], [demo_blocks])
     
     for i, btn in enumerate(tokens_container):
         btn.click(partial(run_interpretation, i=i), [interpretation_prompt,