Update app.py

app.py

@@ -13,6 +13,7 @@ from transformers import PreTrainedModel, PreTrainedTokenizer, AutoModelForCausa
 from interpret import InterpretationPrompt
 
 MAX_PROMPT_TOKENS = 60
+MAX_NUM_LAYERS = 50
 
 
 ## info
@@ -102,7 +103,7 @@ def get_hidden_states(raw_original_prompt):
     token_btns = ([gr.Button(token, visible=True) for token in tokens] 
                   + [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(len(interpretation_bubbles))]
+    invisible_bubbles = [gr.Textbox('', visible=False) for i in range(MAX_NUM_LAYERS)]
     global_state.hidden_states = hidden_states
     return [progress_dummy_output, *token_btns, *invisible_bubbles]
 
@@ -136,9 +137,9 @@ def run_interpretation(raw_interpretation_prompt, max_new_tokens, do_sample,
     generated = interpretation_prompt.generate(global_state.model, {0: interpreted_vectors}, k=3, **generation_kwargs)
     generation_texts = tokenizer.batch_decode(generated)
     progress_dummy_output = ''
-    return ([progress_dummy_output] + 
-            [gr.Textbox(text.replace('\n', ' '), visible=True, container=False, label=f'Layer {i}') for text in generation_texts]
-           )
+    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))]
+    return [progress_dummy_output, *bubble_outputs]
 
 
 ## main
@@ -235,7 +236,7 @@ with gr.Blocks(theme=gr.themes.Default(), css='styles.css') as demo:
     progress_dummy = gr.Markdown('', elem_id='progress_dummy')
     interpretation_bubbles = [gr.Textbox('', container=False, visible=False, 
                                          elem_classes=['bubble', 'even_bubble' if i % 2 == 0 else 'odd_bubble']
-                                        ) for i in range(model.config.num_hidden_layers)]
+                                        ) for i in range(MAX_NUM_LAYERS)]
 
     
     # event listeners