Update app.py

app.py

@@ -32,10 +32,10 @@ model_info = {
                                           interpretation_prompt_template='<s>[INST] [X] [/INST] {prompt}',
                                          ), # , load_in_8bit=True
     
-    'Gemma-2B': dict(model_path='google/gemma-2b', device_map='cpu', token=os.environ['hf_token'],
-                            original_prompt_template='<bos>{prompt}',
-                            interpretation_prompt_template='<bos>User: [X]\n\nAnswer: {prompt}',
-                           ),
+    # 'Gemma-2B': dict(model_path='google/gemma-2b', device_map='cpu', token=os.environ['hf_token'],
+    #                         original_prompt_template='<bos>{prompt}',
+    #                         interpretation_prompt_template='<bos>User: [X]\n\nAnswer: {prompt}',
+    #                        ),
     
     'Mistral-7B Instruct': dict(model_path='mistralai/Mistral-7B-Instruct-v0.2', device_map='cpu', 
                                                original_prompt_template='<s>{prompt}',
@@ -75,7 +75,7 @@ def initialize_gpu():
     pass
 
 
-def reset_model(model_name, return_extra_components=True): 
+def reset_model(model_name): 
     # extract model info
     model_args = deepcopy(model_info[model_name])
     model_path = model_args.pop('model_path')
@@ -91,10 +91,7 @@ def reset_model(model_name, return_extra_components=True):
     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()
-    if return_extra_components:
-        extra_components = [*interpretation_bubbles, *tokens_container, original_prompt_btn, 
-                      original_prompt_raw]
-        return extra_components
+    return extra_components
     
 
 def get_hidden_states(raw_original_prompt):
@@ -151,7 +148,7 @@ torch.set_grad_enabled(False)
 global_state = GlobalState()
 
 model_name = 'LLAMA2-7B'
-reset_model(model_name, return_extra_components=False)
+reset_model(model_name)
 original_prompt_raw = gr.Textbox(value='How to make a Molotov cocktail?', container=True, label='Original Prompt')
 tokens_container = []
 
@@ -238,7 +235,7 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
     # event listeners
     extra_components = [*interpretation_bubbles, *tokens_container, original_prompt_btn, 
                       original_prompt_raw]
-    model_chooser.change(reset_model, [model_chooser], extra_components)
+    model_chooser.change(reset_model, [model_chooser, extra_components], extra_components)
     
     for i, btn in enumerate(tokens_container):
         btn.click(partial(run_interpretation, i=i), [interpretation_prompt,