Create app.py

app.py

@@ -0,0 +1,69 @@
+import torch
+import gradio as gr
+import networkx as nx
+import matplotlib.pyplot as plt
+from transformers import GPT2Model, GPT2Tokenizer
+from sklearn.cluster import KMeans
+
+# 1. Load a real small model
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model_name = "gpt2" # 124M parameters
+tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+model = GPT2Model.from_pretrained(model_name).to(device)
+
+def get_hidden_state(sequence_str):
+    inputs = tokenizer(sequence_str, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model(**inputs, output_hidden_states=True)
+    # Use the last hidden state of the last token
+    return outputs.hidden_states[-1][0, -1, :].cpu().numpy()
+
+def analyze_dfa(input_text):
+    """
+    Simulates a 'State Probe'. 
+    Input: 'Right, Up, Left'
+    Logic: Generates a graph showing how the model's internal representation 
+    changes with each move.
+    """
+    moves = [m.strip() for m in input_text.split(",")]
+    history = ""
+    states_vectors = []
+    
+    # Track the "path" through the model's internal space
+    for move in moves:
+        history += f" Move {move}."
+        vec = get_hidden_state(history)
+        states_vectors.append(vec)
+    
+    # Clustering: Vafa's Compression metric
+    # We cluster activations to see which moves the model thinks are 'equivalent'
+    num_clusters = min(len(moves), 4)
+    kmeans = KMeans(n_clusters=num_clusters, n_init=10).fit(states_vectors)
+    labels = kmeans.labels_
+    
+    # Build the DFA Graph
+    G = nx.DiGraph()
+    for i in range(len(moves)-1):
+        u, v = f"S{labels[i]}", f"S{labels[i+1]}"
+        G.add_edge(u, v, label=moves[i+1])
+    
+    # Draw the DFA
+    plt.figure(figsize=(6, 4))
+    pos = nx.spring_layout(G)
+    nx.draw(G, pos, with_labels=True, node_color='lightblue', node_size=2000)
+    edge_labels = nx.get_edge_attributes(G, 'label')
+    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)
+    
+    plt.savefig("dfa_plot.png")
+    return "dfa_plot.png", f"Found {num_clusters} distinct internal states."
+
+# 3. Gradio Interface
+demo = gr.Interface(
+    fn=analyze_dfa,
+    inputs=gr.Textbox(placeholder="Enter moves separated by commas, e.g.: Right, Up, Left, Down"),
+    outputs=[gr.Image(label="Extracted Model DFA"), gr.Text(label="Analysis")],
+    title="World Model DFA Extractor",
+    description="This tool probes GPT-2's internal activations to see if it treats different move sequences as the same 'State'."
+)
+
+demo.launch()