app.py trying to fix lists functionality

app.py

@@ -341,12 +341,13 @@ def initialize_resources():
     
     logger.info("Resources initialized successfully.")
 
+
 @spaces.GPU
-def process_input(selected_sae, feature_number, weight_type, use_token_centroid, scaling_factor, use_pca, pca_weight, num_exp, denom_exp, mode):
+def process_input(selected_sae, feature_number, weight_type, use_token_centroid, scaling_factor, use_pca, pca_weight, num_exp, denom_exp, mode, top_500=False):
     global w_enc_dict, w_dec_dict, model, tokenizer, token_embeddings
     
     try:
-        logger.info("Processing input: SAE={}, feature_number={}, mode={}".format(selected_sae, feature_number, mode))
+        logger.info(f"Processing input: SAE={selected_sae}, feature_number={feature_number}, mode={mode}")
 
         # Load the SAE weights if they are not already loaded
         if selected_sae not in w_enc_dict or selected_sae not in w_dec_dict:
@@ -373,16 +374,20 @@ def process_input(selected_sae, feature_number, weight_type, use_token_centroid,
                 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])
-            
-            logger.info("Returning result for cosine distance token list")
-            return result, None  # Return the result and no image
+            if top_500:
+                # Generate the top 500 list
+                result = "Top 500 list:\n"
+                result += "\n".join([f"{token!r}: {value:.4f}" for token, value in closest_tokens_with_values])
+                logger.info("Returning top 500 list")
+                return result, None
+            else:
+                # Generate the top 100 list
+                token_list = [token for token, _ in closest_tokens_with_values[:100]]
+                result = f"100 tokens whose embeddings produce the smallest ratio:\n\n"
+                result += f"[{', '.join(repr(token) for token in token_list)}]\n"
+                logger.info("Returning top 100 tokens")
+                return result, None
         
-        # In case the mode is not "cosine distance token lists", return a default message
         return "Mode not recognized or not implemented in this step.", None
     
     except Exception as e:
@@ -391,6 +396,8 @@ def process_input(selected_sae, feature_number, weight_type, use_token_centroid,
 
 
 
+
+
 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))
@@ -398,80 +405,28 @@ def trim_tree(trim_cutoff, tree_data):
 
 
 
+
 def gradio_interface():
     def update_visibility(mode):
         if mode == "definition tree generation":
             return gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
         else:
-            return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+            return gr.update(visible=False), gr.update(visible=False), gr.update(visible(False))
 
     def update_neuronpedia(selected_sae, feature_number):
         layer_number = int(selected_sae.split()[-1])
         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
-        
-        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
-
-
+        # Call process_input without generating the top 500 list initially
+        return process_input(selected_sae, feature_number, weight_type, use_token_centroid, scaling_factor, use_pca, pca_weight, num_exp, denom_exp, mode, top_500=False)
 
+    @spaces.GPU
+    def generate_top_500(selected_sae, feature_number, weight_type, use_token_centroid, scaling_factor, use_pca, pca_weight, num_exp, denom_exp, mode):
+        # Call process_input with top_500=True to generate the full list
+        return process_input(selected_sae, feature_number, weight_type, use_token_centroid, scaling_factor, use_pca, pca_weight, num_exp, denom_exp, mode, top_500=True)
 
     def trim_tree(trim_cutoff, tree_data):
         if tree_data is None:
@@ -481,8 +436,8 @@ def gradio_interface():
         return trimmed_tree_image
 
     with gr.Blocks() as demo:
-        gr.Markdown("# Gemma-2B SAE Feature Explorer (gradual3)")
-        
+        gr.Markdown("# Gemma-2B SAE Feature Explorer (almost there?)")
+
         with gr.Row():
             with gr.Column(scale=2):
                 selected_sae = gr.Dropdown(choices=["Gemma-2B layer 0", "Gemma-2B layer 6", "Gemma-2B layer 10", "Gemma-2B layer 12"], label="Select SAE")
@@ -508,6 +463,9 @@ def gradio_interface():
                 output_text = gr.Textbox(label="Output", lines=20)
                 output_image = gr.Image(label="Tree Diagram", visible=False)
                 
+                generate_top_500_btn = gr.Button("Generate Top 500 Tokens and Power Ratios", visible=False)
+                output_500_text = gr.Textbox(label="Top 500 Output", lines=20, visible=False)
+
                 trim_slider = gr.Slider(minimum=0.00001, maximum=0.1, value=0.00001, label="Trim cutoff for cumulative probability", visible=False)
                 trim_btn = gr.Button("Trim Tree", visible=False)
 
@@ -515,28 +473,40 @@ def gradio_interface():
         neuronpedia_html = gr.HTML(label="Neuronpedia")
 
         inputs = [selected_sae, feature_number, weight_type, use_token_centroid, scaling_factor, use_pca, pca_weight, num_exp, denom_exp, mode]
-        
 
         generate_btn.click(
-            process_input,
+            update_output,
             inputs=inputs,
-            outputs=[output_text, output_image],  # Ensure these match your components
+            outputs=[output_text, output_image],
             show_progress="full"
         )
 
-
-        trim_btn.click(trim_tree, inputs=[trim_slider, tree_data_state], outputs=[output_image])
+        generate_top_500_btn.click(
+            generate_top_500,
+            inputs=inputs,
+            outputs=[output_500_text],
+            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])
 
-
+        output_text.change(
+            lambda text: (gr.update(visible=True), gr.update(visible=True)) if "100 tokens" in text else (gr.update(visible(False)), gr.update(visible(False))),
+            inputs=[output_text],
+            outputs=[generate_top_500_btn, output_500_text]
+        )
 
     return demo
 
 
+
+
+
 if __name__ == "__main__":
     try:
         logger.info("Starting application initialization...")