Update app.py

app.py

@@ -6,35 +6,53 @@ from sklearn.cluster import KMeans
 import networkx as nx
 import matplotlib.pyplot as plt
 
-# 1. Configuration for Models & Datasets
+# 1. Configuration for Models & Specific Dataset Configs
 MODELS = ["gpt2", "distilgpt2", "qwen/Qwen2.5-0.5B", "TinyLlama/TinyLlama-1.1B-Chat-v1.0"]
-DATASETS = ["wikitext", "tinystories", "ag_news"]
 
-def analyze_world_model(model_name, dataset_name, num_samples=20):
-    # Load Model & Tokenizer
+# Updated to include the specific config names required by HuggingFace
+DATASET_CONFIGS = {
+    "wikitext (v2-raw)": ("wikitext", "wikitext-2-raw-v1"),
+    "wikitext (v103-raw)": ("wikitext", "wikitext-103-raw-v1"),
+    "TinyStories": ("roneneldan/TinyStories", None),
+    "AG News": ("ag_news", None)
+}
+
+def analyze_world_model(model_name, dataset_key, num_samples=20):
     device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    # Get the dataset name and its config
+    dataset_name, config_name = DATASET_CONFIGS[dataset_key]
+    
+    # Load Model & Tokenizer
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModel.from_pretrained(model_name).to(device)
     
-    # Load Dataset
-    ds = load_dataset(dataset_name, split='train', streaming=True).take(num_samples)
+    # 2. FIXED: Load Dataset with config_name
+    try:
+        if config_name:
+            # Passes both dataset name and the specific config
+            ds = load_dataset(dataset_name, config_name, split='train', streaming=True).take(num_samples)
+        else:
+            ds = load_dataset(dataset_name, split='train', streaming=True).take(num_samples)
+    except Exception as e:
+        return None, f"Error loading dataset: {str(e)}"
     
     all_hidden_states = []
-    labels = []
     
-    # Step A: The Probe (Keplerian Observation)
+    # Step A: The Probe (Hidden State Extraction)
     for i, example in enumerate(ds):
-        text = example['text'][:100] # Use a snippet
+        # Handle different dataset structures (some use 'text', some use 'content')
+        text = example.get('text', example.get('content', ''))[:100]
+        if not text: continue
+            
         inputs = tokenizer(text, return_tensors="pt").to(device)
         with torch.no_grad():
             outputs = model(**inputs, output_hidden_states=True)
-            # Take the last hidden state of the sequence
+            # Snapshot of the last layer's representation
             state = outputs.hidden_states[-1][0, -1, :].cpu().numpy()
             all_hidden_states.append(state)
-            labels.append(f"Seq_{i}")
 
-    # Step B: Myhill-Nerode Clustering (Newtonian Recovery)
-    # We cluster to find 'Equivalence Classes' (Internal States)
+    # Step B: Newtonian Recovery (Clustering)
     n_clusters = min(len(all_hidden_states), 5)
     kmeans = KMeans(n_clusters=n_clusters, n_init=10).fit(all_hidden_states)
     state_assignments = kmeans.labels_
@@ -43,25 +61,30 @@ def analyze_world_model(model_name, dataset_name, num_samples=20):
     G = nx.DiGraph()
     for i in range(len(state_assignments) - 1):
         u, v = f"S{state_assignments[i]}", f"S{state_assignments[i+1]}"
-        G.add_edge(u, v, label=f"Next_{i}")
+        G.add_edge(u, v)
 
-    # Draw the DFA
     plt.figure(figsize=(8, 6))
     pos = nx.spring_layout(G)
-    nx.draw(G, pos, with_labels=True, node_color='orange', node_size=3000, font_weight='bold')
+    nx.draw(G, pos, with_labels=True, node_color='orange', node_size=3000, font_weight='bold', arrowsize=20)
     plt.savefig("dfa_output.png")
+    plt.close() # Clean up memory
     
     return "dfa_output.png", f"Model '{model_name}' reduced this dataset into {n_clusters} distinct internal states."
 
 # 3. Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("# The Universal Newtonian Probe")
+    gr.Markdown("Analyze how models build internal maps of different datasets.")
+    
     with gr.Row():
         m_drop = gr.Dropdown(choices=MODELS, label="Select Model", value="gpt2")
-        d_drop = gr.Dropdown(choices=DATASETS, label="Select Dataset", value="wikitext")
+        d_drop = gr.Dropdown(choices=list(DATASET_CONFIGS.keys()), label="Select Dataset", value="wikitext (v2-raw)")
+    
     btn = gr.Button("Analyze Coherence")
-    out_img = gr.Image(label="Extracted DFA")
-    out_txt = gr.Textbox(label="Analysis Result")
+    
+    with gr.Row():
+        out_img = gr.Image(label="Extracted DFA")
+        out_txt = gr.Textbox(label="Analysis Result")
     
     btn.click(analyze_world_model, inputs=[m_drop, d_drop], outputs=[out_img, out_txt])