app.py refactored

app.py

@@ -420,6 +420,7 @@ def gradio_interface():
         url = get_neuronpedia_url(layer_number, feature_number)
         return f'<iframe src="{url}" width="100%" height="300px"></iframe>'
 
+
     @spaces.GPU
     def update_output(selected_sae, feature_number, weight_type, use_token_centroid, scaling_factor, use_pca, pca_weight, num_exp, denom_exp, mode):
         global w_enc_dict, w_dec_dict, model, tokenizer, token_embeddings
@@ -479,6 +480,9 @@ def gradio_interface():
             logger.error(traceback.format_exc())
             return f"Error: {str(e)}\nPlease check the logs for more details.", None, None
 
+
+
+
     def trim_tree(trim_cutoff, tree_data):
         if tree_data is None:
             return None
@@ -522,14 +526,24 @@ def gradio_interface():
 
         inputs = [selected_sae, feature_number, weight_type, use_token_centroid, scaling_factor, use_pca, pca_weight, num_exp, denom_exp, mode]
         
-        generate_btn.click(update_output, inputs=inputs, outputs=[output_text, output_image, tree_data_state])
+
+        generate_btn.click(
+            update_output,
+            inputs=inputs,
+            outputs=[output_text, output_image, tree_data_state],
+            show_progress="full"
+        )
+
+
         trim_btn.click(trim_tree, inputs=[trim_slider, tree_data_state], outputs=[output_image])
         
         mode.change(update_visibility, inputs=[mode], outputs=[output_image, trim_slider, trim_btn])
-        
+            
         selected_sae.change(update_neuronpedia, inputs=[selected_sae, feature_number], outputs=[neuronpedia_html])
         feature_number.change(update_neuronpedia, inputs=[selected_sae, feature_number], outputs=[neuronpedia_html])
-    
+
+
+
     return demo