app.py refactored

app.py

@@ -403,6 +403,65 @@ def process_input(selected_sae, feature_number, weight_type, use_token_centroid,
 
 
 def gradio_interface():
+    @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
+        
+        try:
+            if selected_sae not in w_enc_dict or selected_sae not in w_dec_dict:
+                w_enc, w_dec = load_sae_weights(selected_sae)
+                if w_enc is None or w_dec is None:
+                    return f"Failed to load SAE weights for {selected_sae}. Please try a different SAE or check your connection.", None, None
+                w_enc_dict[selected_sae] = w_enc
+                w_dec_dict[selected_sae] = w_dec
+            else:
+                w_enc, w_dec = w_enc_dict[selected_sae], w_dec_dict[selected_sae]
+            
+            token_centroid = torch.mean(token_embeddings, dim=0)
+            feature_vector = create_feature_vector(w_enc, w_dec, int(feature_number), weight_type, token_centroid, use_token_centroid, scaling_factor)
+            
+            if use_pca:
+                pca_direction = perform_pca(token_embeddings)
+                feature_vector = create_ghost_token(feature_vector, token_centroid, pca_direction, scaling_factor, pca_weight)
+            
+            if mode == "cosine distance token lists":
+                closest_tokens_with_values = find_closest_tokens(
+                    feature_vector, token_embeddings, tokenizer, 
+                    top_k=500, num_exp=num_exp, denom_exp=denom_exp
+                )
+                
+                token_list = [token for token, _ in closest_tokens_with_values]
+                result = f"100 tokens whose embeddings produce the smallest ratio:\n\n"
+                result += f"[{', '.join(repr(token) for token in token_list[:100])}]\n\n"
+                result += "Top 500 list:\n"
+                result += "\n".join([f"{token!r}: {value:.4f}" for token, value in closest_tokens_with_values])
+                
+                return result, None, None  # Return the result, no image, and no tree data
+            elif mode == "definition tree generation":
+                base_prompt = f'A typical definition of "{tokenizer.decode([config.SUB_TOKEN_ID], skip_special_tokens=True)}" would be "'
+                tree_generator = generate_definition_tree(base_prompt, feature_vector, model, tokenizer, config)
+                
+                # Generate the tree
+                tree_text = ""
+                tree_dict = None
+                for item in tree_generator:
+                    if isinstance(item, str):
+                        tree_text += item
+                        yield tree_text, None, None  # Yield the updated text, no image, and no tree data
+                    else:
+                        tree_dict = item
+                
+                # Generate the tree visualization
+                max_weight, min_weight = find_max_min_cumulative_weight(tree_dict)
+                tree_image = create_tree_diagram(tree_dict, config, max_weight, min_weight)
+                
+                return tree_text, tree_image, tree_dict  # Return the final text, tree image, and tree data
+            
+        except Exception as e:
+            logger.error(f"Error in update_output: {str(e)}")
+            logger.error(traceback.format_exc())
+            return f"Error: {str(e)}\nPlease check the logs for more details.", None, None
+
     with gr.Blocks() as demo:
         gr.Markdown("# Gemma-2B SAE Feature Explorer")
         
@@ -434,34 +493,12 @@ def gradio_interface():
                 trim_slider = gr.Slider(minimum=0.00001, maximum=0.1, value=0.00001, label="Trim cutoff for cumulative probability")
                 trim_btn = gr.Button("Trim Tree")
 
-        def update_output(selected_sae, feature_number, weight_type, use_token_centroid, scaling_factor, use_pca, pca_weight, num_exp, denom_exp, mode):
-            result = process_input(selected_sae, feature_number, weight_type, use_token_centroid, scaling_factor, use_pca, pca_weight, num_exp, denom_exp, mode)
-    
-            if mode == "definition tree generation":
-                tree_text = ""
-                tree_image = None
-                tree_data = None
-                for i, item in enumerate(result):
-                    progress(i / 100)  # Assuming max 100 iterations, adjust as needed
-                    if isinstance(item, str):
-                        tree_text += item
-                        yield tree_text, tree_image, tree_data
-                    else:
-                        tree_data = item
-                        tree_image = create_tree_diagram(tree_data, config, *find_max_min_cumulative_weight(tree_data))
-                yield tree_text, tree_image, tree_data
-            else:
-                yield result, None, None
-
-        def trim_tree(trim_cutoff, tree_data):
-            max_weight, min_weight = find_max_min_cumulative_weight(tree_data)
-            trimmed_tree_image = create_tree_diagram(tree_data, config, max_weight, min_weight, trim_cutoff=float(trim_cutoff))
-            return trimmed_tree_image
+        tree_data_state = gr.State()
 
         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, gr.State(key="tree_data")])
-        trim_btn.click(trim_tree, inputs=[trim_slider, gr.State(key="tree_data")], outputs=[output_image])
+        generate_btn.click(update_output, inputs=inputs, outputs=[output_text, output_image, tree_data_state])
+        trim_btn.click(trim_tree, inputs=[trim_slider, tree_data_state], outputs=[output_image])
     
     return demo