back to last promising version with treebuild crashlog trying to move text to interface with 4o help AGAIN2

app.py

@@ -187,6 +187,11 @@ def update_token_embedding(model, token_id, new_embedding):
     new_embedding = new_embedding.to(model.get_input_embeddings().weight.device)
     model.get_input_embeddings().weight.data[token_id] = new_embedding
 
+@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, progress=gr.Progress()):
+    # Call process_input to generate the output
+    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, progress=progress)
+
 def produce_next_token_ids(input_ids, model, topk, sub_token_id):
     input_ids = input_ids.to(model.device)
     with torch.no_grad():
@@ -445,23 +450,30 @@ 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), gr.update(visible=False), gr.update(visible=False)
+            return (gr.update(visible=True), gr.update(visible=True), gr.update(visible=True),
+                    gr.update(visible=False), gr.update(visible=False))
         else:
-            return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)
+            return (gr.update(visible=False), gr.update(visible=False), gr.update(visible=False),
+                    gr.update(visible=True), gr.update(visible=True))
+
+    @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, progress=gr.Progress()):
+        # Call process_input to generate the output
+        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, progress=progress)
 
     with gr.Blocks() as demo:
         gr.Markdown("# Gemma-2B SAE Feature Explorer")
 
         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")
+                selected_sae = gr.Dropdown(choices=["Gemma-2B layer 0", "Gemma-2B layer 6", "Gemma-2B layer 10", "Gemma-2B layer 12"],
+                                           label="Select SAE")
                 feature_number = gr.Number(label="Select feature number", minimum=0, maximum=16383, value=0)
                 
-                mode = gr.Radio(
-                    choices=["cosine distance token lists", "definition tree generation"],
-                    label="Select mode",
-                    value="cosine distance token lists"
-                )
+                mode = gr.Radio(choices=["cosine distance token lists", "definition tree generation"],
+                                label="Select mode", value="cosine distance token lists")
                 
                 weight_type = gr.Radio(["encoder", "decoder"], label="Select weight type for feature vector construction", value="encoder")
                 use_token_centroid = gr.Checkbox(label="Use token centroid offset", value=True)
@@ -485,16 +497,12 @@ def gradio_interface():
 
         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]
+        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],  # Now the text output will be displayed in the Textbox
-            show_progress="full"
-        )
+        generate_btn.click(update_output, inputs=inputs, outputs=[output_text, output_image], show_progress="full")
 
-        # other buttons and changes remain the same...
+        # Add other button functionality as needed...
 
     return demo
 
@@ -502,7 +510,6 @@ def gradio_interface():
 
 
 
-
 if __name__ == "__main__":
     try:
         logger.info("Starting application initialization...")