Update app.py

app.py

@@ -41,7 +41,7 @@ suggested_interpretation_prompts = [
 def initialize_gpu():
     pass
 
-def reset_model(model_name, *extra_components): 
+def reset_model(model_name, *extra_components, with_extra_components=True): 
     # extract model info
     model_args = deepcopy(model_info[model_name])
     model_path = model_args.pop('model_path')
@@ -58,7 +58,10 @@ def reset_model(model_name, *extra_components):
     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 extra_components
+    if with_extra_components:
+        for x in interpretation_bubbles:
+            x.visible = False
+        return extra_components
     
 
 def get_hidden_states(raw_original_prompt):
@@ -101,8 +104,8 @@ def run_interpretation(raw_interpretation_prompt, max_new_tokens, do_sample,
     interpretation_prompt = InterpretationPrompt(global_state.tokenizer, interpretation_prompt)
 
     # 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}, layers_format=global_state.layers_format, k=3, 
+    generated = interpretation_prompt.generate(global_state.model, {0: interpreted_vectors}, 
+                                               layers_format=global_state.layers_format, k=3, 
                                                **generation_kwargs)
     generation_texts = global_state.tokenizer.batch_decode(generated)
     progress_dummy_output = ''
@@ -116,7 +119,7 @@ torch.set_grad_enabled(False)
 global_state = GlobalState()
 
 model_name = 'LLAMA2-7B'
-reset_model(model_name)
+reset_model(model_name, with_extra_components=False)
 original_prompt_raw = gr.Textbox(value='How to make a Molotov cocktail?', container=True, label='Original Prompt')
 tokens_container = []