Update app.py

app.py

@@ -42,7 +42,7 @@ def analyze_dfa(input_text):
         vec = get_hidden_state(history)
         states_vectors.append(vec)
     
-    # --- 1. KMeans Graph (Unsupervised State Map) ---
+    # --- 1. KMeans Graph (Discrete State Machine) ---
     num_clusters = min(len(moves), 4)
     kmeans = KMeans(n_clusters=num_clusters, n_init=10).fit(states_vectors)
     km_labels = kmeans.labels_
@@ -51,40 +51,68 @@ def analyze_dfa(input_text):
     for i in range(len(moves)-1):
         G_km.add_edge(f"S{km_labels[i]}", f"S{km_labels[i+1]}", label=moves[i+1])
     
-    plt.figure(figsize=(12, 5))
-    plt.subplot(1, 2, 1)
+    plt.figure(figsize=(8, 6))
     pos_km = nx.spring_layout(G_km)
-    nx.draw(G_km, pos_km, with_labels=True, node_color='lightblue', node_size=1500)
-    nx.draw_networkx_edge_labels(G_km, pos_km, edge_labels=nx.get_edge_attributes(G_km, 'label'))
-    plt.title("KMeans DFA (State-Based)")
+    nx.draw(G_km, pos_km, with_labels=True, node_color='lightblue', node_size=2500, font_size=12, font_weight='bold')
+    nx.draw_networkx_edge_labels(G_km, pos_km, edge_labels=nx.get_edge_attributes(G_km, 'label'), font_size=10)
+    plt.title("Logical State Machine (KMeans)")
+    km_plot = "km_plot.png"
+    plt.savefig(km_plot, dpi=150)
+    plt.close()
 
-    # --- 2. Linear Probe / PCA (Geometric Map) ---
-    logger.info("📐 Running Linear Probe (PCA) to find the 'Spatial Axis'...")
+    # --- 2. Linear Probe PCA (Geometric State Machine) ---
     pca = PCA(n_components=2)
     coords = pca.fit_transform(states_vectors)
     
-    plt.subplot(1, 2, 2)
-    plt.scatter(coords[:, 0], coords[:, 1], c=range(len(moves)), cmap='viridis', s=100)
+    plt.figure(figsize=(10, 8)) # Increased size for better visibility
+    plt.scatter(coords[:, 0], coords[:, 1], c=range(len(moves)), cmap='viridis', s=200, edgecolors='black')
+    
+    # Drawing arrows between coordinates (The Linear Probe "State Machine")
+    for i in range(len(moves)-1):
+        plt.arrow(coords[i, 0], coords[i, 1], 
+                  coords[i+1, 0] - coords[i, 0], 
+                  coords[i+1, 1] - coords[i, 1], 
+                  head_width=5, length_includes_head=True, alpha=0.5, color='gray')
+    
     for i, move in enumerate(moves):
-        plt.annotate(f"{i}:{move}", (coords[i, 0], coords[i, 1]))
-    plt.plot(coords[:, 0], coords[:, 1], 'r--', alpha=0.3) # Path line
-    plt.title("Linear Probe (Spatial Projection)")
+        plt.annotate(f"Step {i}: {move}", (coords[i, 0], coords[i, 1]), 
+                     xytext=(5, 5), textcoords='offset points', fontsize=9, fontweight='bold')
+    
+    plt.grid(True, linestyle='--', alpha=0.6)
+    plt.title("Geometric State Machine (Linear Probe PCA)")
+    plt.xlabel("Principal Component 1 (Primary Axis of Variance)")
+    plt.ylabel("Principal Component 2 (Secondary Axis of Variance)")
     
-    plot_path = "comparison_plot.png"
-    plt.savefig(plot_path)
+    pca_plot = "pca_plot.png"
+    plt.savefig(pca_plot, dpi=150)
     plt.close()
     
-    return plot_path, f"KMeans Labels: {km_labels}", log_capture.getvalue()
+    return km_plot, pca_plot, f"Labels: {km_labels}", log_capture.getvalue()
 
-# Launching with dual display
-with gr.Blocks() as demo:
-    gr.Markdown("# KMeans vs. Linear Probe Analysis")
-    input_box = gr.Textbox(label="Moves (Right, Left...)")
-    submit_btn = gr.Button("Compare")
+# Gradio Interface with Separated Columns
+with gr.Blocks(title="World Model Hybrid Probe") as demo:
+    gr.Markdown("# 🛰️ World Model Hybrid Probe")
+    gr.Markdown("Comparing **Logical Categorization** (KMeans) vs **Spatial Intuition** (Linear PCA).")
+    
+    with gr.Row():
+        input_box = gr.Textbox(label="Input Moves", placeholder="Up, Up, Right, Left", scale=4)
+        submit_btn = gr.Button("Analyze", variant="primary", scale=1)
+    
     with gr.Row():
-        output_img = gr.Image(label="KMeans (Left) vs Linear PCA (Right)")
-        analysis_text = gr.Textbox(label="Mapping Results")
-    log_box = gr.Textbox(label="Probe Logs", lines=5)
-    submit_btn.click(analyze_dfa, input_box, [output_img, analysis_text, log_box])
+        # Box 1: Logic
+        with gr.Column(variant="panel"):
+            gr.Markdown("### 1. Discrete State Logic (DFA)")
+            output_km = gr.Image(label="KMeans DFA", type="filepath")
+            analysis_text = gr.Textbox(label="Cluster Labels", interactive=False)
+        
+        # Box 2: Geometry (The Clearer Linear Probe)
+        with gr.Column(variant="panel"):
+            gr.Markdown("### 2. Geometric Trajectory (Linear Probe)")
+            output_pca = gr.Image(label="Spatial PCA Map", type="filepath")
+            gr.Markdown("*This map shows the 'Mental Path' GPT-2 takes through its vector space.*")
+
+    log_box = gr.Textbox(label="Probe Logs", lines=5, interactive=False)
+
+    submit_btn.click(analyze_dfa, input_box, [output_km, output_pca, analysis_text, log_box])
 
 demo.launch()