Update app.py

app.py

@@ -75,7 +75,7 @@ def initialize_gpu():
     pass
 
 
-def reset_model(model_name, *extra_components): 
+def reset_model(model_name): 
     # extract model info
     model_args = deepcopy(model_info[model_name])
     model_path = model_args.pop('model_path')
@@ -113,7 +113,7 @@ 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]
     length_penalty = -length_penalty   # unintuitively, length_penalty > 0 will make sequences longer, so we negate it
 
@@ -146,6 +146,7 @@ def run_interpretation(raw_interpretation_prompt, max_new_tokens, do_sample,
 ## main
 torch.set_grad_enabled(False)
 global_state = GlobalState()
+extra_components = []
 
 model_name = 'LLAMA2-7B'
 reset_model(model_name)
@@ -183,7 +184,7 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
         #     gr.Markdown('<span style="font-size:180px;">🤔</span>')
 
     with gr.Group():
-        model_chooser = gr.Radio(choices=list(model_info.keys()), value=model_name)
+        model_chooser = gr.Radio(label='Model', choices=list(model_info.keys()), value=model_name)
     
     with gr.Blocks() as demo_blocks:
         gr.Markdown('## Choose Your Interpretation Prompt')
@@ -233,9 +234,9 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
     
     
     # event listeners
-    all_components = [*interpretation_bubbles, *tokens_container, original_prompt_btn, 
+    extra_components = [*interpretation_bubbles, *tokens_container, original_prompt_btn, 
                       original_prompt_raw]
-    model_chooser.change(reset_model, [model_chooser, *all_components], all_components)
+    model_chooser.change(reset_model, [model_chooser], extra_components)
     
     for i, btn in enumerate(tokens_container):
         btn.click(partial(run_interpretation, i=i), [interpretation_prompt,