Update app.py

app.py

@@ -34,6 +34,16 @@ for name, layer in model.named_modules():
     if 'layer' in name or 'embeddings' in name:
         layer.register_forward_hook(lambda m, i, o, n=name: hook_fn(m, i, o, n))
 
+def convert_dict_keys_to_str(d):
+    """Recursively convert all dictionary keys to strings."""
+    if isinstance(d, dict):
+        return {str(k): convert_dict_keys_to_str(v) for k, v in d.items()}
+    elif isinstance(d, list):
+        return [convert_dict_keys_to_str(item) for item in d]
+    elif isinstance(d, np.ndarray):
+        return d.tolist()  # Convert numpy arrays to lists
+    return d
+
 def process_input(input_text, layer_name, visualize_option, attribution_target=0):
     """
     Process input text, compute embeddings, activations, attention, and attribution.
@@ -44,7 +54,7 @@ def process_input(input_text, layer_name, visualize_option, attribution_target=0
     - attribution_target: Target class for attribution (0 or 1)
     Returns:
     - HTML string with base64-encoded image(s)
-    - List of dataframe dictionaries
+    - List of dataframe dictionaries with string keys
     - Status message
     """
     global activations
@@ -94,10 +104,11 @@ def process_input(input_text, layer_name, visualize_option, attribution_target=0
                     # Dataframe for coordinates
                     dataframe = pd.DataFrame({
                         "Token": tokens,
-                        "t-SNE X": reduced[:, 0],
-                        "t-SNE Y": reduced[:, 1]
-                    }).to_dict()
-                    dataframes.append(dataframe)
+                        "t-SNE_X": reduced[:, 0],
+                        "t-SNE_Y": reduced[:, 1]
+                    })
+                    dataframe.index = [f"idx_{i}" for i in range(len(dataframe))]  # String indices
+                    dataframes.append(convert_dict_keys_to_str(dataframe.to_dict()))
                 except Exception as e:
                     logger.warning(f"t-SNE failed: {e}")
                     dataframes.append({"Error": [str(e)]})
@@ -122,8 +133,8 @@ def process_input(input_text, layer_name, visualize_option, attribution_target=0
                 html_plots.append(f'<img src="data:image/png;base64,{img_base64}" alt="Attention Heatmap" style="max-width:100%;"/>')
                 plt.close()
                 # Dataframe for attention weights
-                dataframe = pd.DataFrame(attn, index=tokens, columns=tokens).to_dict()
-                dataframes.append(dataframe)
+                dataframe = pd.DataFrame(attn, index=tokens, columns=[f"token_{i}" for i in range(len(tokens))])
+                dataframes.append(convert_dict_keys_to_str(dataframe.to_dict()))
             else:
                 dataframes.append({"Error": ["No attention weights available."]})
                 html_plots.append("<p>Error: No attention weights available.</p>")
@@ -135,8 +146,8 @@ def process_input(input_text, layer_name, visualize_option, attribution_target=0
                 if isinstance(act, tuple):
                     act = act[0]
                 act = act[0].detach().numpy()  # [seq_len, hidden_size]
-                dataframe = pd.DataFrame(act, index=tokens).to_dict()
-                dataframes.append(dataframe)
+                dataframe = pd.DataFrame(act, index=tokens, columns=[f"dim_{i}" for i in range(act.shape[1])])
+                dataframes.append(convert_dict_keys_to_str(dataframe.to_dict()))
                 # Plot mean activation per token
                 fig, ax = plt.subplots(figsize=(8, 6))
                 mean_act = np.mean(act, axis=1)
@@ -168,8 +179,9 @@ def process_input(input_text, layer_name, visualize_option, attribution_target=0
                 return_convergence_delta=True
             )
             attr = attributions[0].detach().numpy().sum(axis=1)
-            attr_df = pd.DataFrame({"Token": tokens, "Attribution": attr}).to_dict()
-            dataframes.append(attr_df)
+            attr_df = pd.DataFrame({"Token": tokens, "Attribution": attr})
+            attr_df.index = [f"idx_{i}" for i in range(len(attr_df))]  # String indices
+            dataframes.append(convert_dict_keys_to_str(attr_df.to_dict()))
             # Plot attributions
             fig, ax = plt.subplots(figsize=(8, 6))
             ax.bar(range(len(attr)), attr)
@@ -180,7 +192,7 @@ def process_input(input_text, layer_name, visualize_option, attribution_target=0
             plt.savefig(buf, format='png', bbox_inches='tight')
             buf.seek(0)
             img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
-            html_plots.append(f'<img src="data:image/png;base64,{img_base64}" alt="Attribution Plot" style="max-width:100%;"/>')
+            html_plots.append(f'<img src="data:image/png;base64,{img_base64}" alt="Attribution Plot" style="max-width:100?%"/>')
             plt.close()
         except Exception as e:
             logger.warning(f"Integrated Gradients failed: {e}")