Update app.py

app.py

@@ -42,8 +42,9 @@ def create_dashboard(feature_id: int, model: Model) -> str:
         </div>
     """
 
-def handle_feature_click(feature_id: int, model: Model):
-    return create_dashboard(feature_id, model)
+def handle_feature_click(feature_id: int, model: str):
+    selected_model = Model.GEMMA if model == "Gemini" else Model.GPT2
+    return create_dashboard(feature_id, selected_model)
 
 def analyze_text(text: str, selected_model: str):
     model = Model.GEMMA if selected_model == "Gemini" else Model.GPT2
@@ -151,7 +152,24 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
             )
         
         with gr.Column(scale=2):
-            features_list = gr.JSON(label="Features", visible=False)
+            @gr.render(inputs=[features_state, current_model])
+            def render_features(features, model):
+                if not features:
+                    return
+                
+                for token_group in features:
+                    gr.Markdown(f"### {token_group['token']}")
+                    with gr.Row():
+                        for feature in token_group['features']:
+                            btn = gr.Button(
+                                f"Feature {feature['id']} (Activation: {feature['activation']:.2f})",
+                                elem_classes=["feature-button"]
+                            )
+                            btn.click(
+                                fn=lambda fid=feature['id']: handle_feature_click(fid, model),
+                                outputs=dashboard
+                            )
+            
             dashboard = gr.HTML()
     
     def update_and_analyze(text, model):