Update app.py

app.py

@@ -5,87 +5,95 @@ from datasets import load_dataset
 from sklearn.cluster import KMeans
 import networkx as nx
 import matplotlib.pyplot as plt
+import collections
 
-# 1. Configuration for Models & Specific Dataset Configs
+# 1. Models & Datasets Configs
 MODELS = ["gpt2", "distilgpt2", "qwen/Qwen2.5-0.5B", "TinyLlama/TinyLlama-1.1B-Chat-v1.0"]
-
-# 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):
+def analyze_world_model(model_name, dataset_key, num_samples=25):
     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)
     
-    # 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)}"
+    # Load Dataset
+    if config_name:
+        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)
     
     all_hidden_states = []
+    input_snippets = []
     
-    # Step A: The Probe (Hidden State Extraction)
+    # Step A: Probe (Hidden State Extraction)
     for i, example in enumerate(ds):
-        # Handle different dataset structures (some use 'text', some use 'content')
-        text = example.get('text', example.get('content', ''))[:100]
+        text = example.get('text', example.get('content', ''))[:150].strip()
         if not text: continue
             
         inputs = tokenizer(text, return_tensors="pt").to(device)
         with torch.no_grad():
             outputs = model(**inputs, output_hidden_states=True)
-            # Snapshot of the last layer's representation
-            state = outputs.hidden_states[-1][0, -1, :].cpu().numpy()
+            # We take the middle-to-late layer where semantic 'World Models' reside
+            state = outputs.hidden_states[-2][0, -1, :].cpu().numpy()
             all_hidden_states.append(state)
+            input_snippets.append(text)
 
     # Step B: Newtonian Recovery (Clustering)
-    n_clusters = min(len(all_hidden_states), 5)
+    n_clusters = 5
     kmeans = KMeans(n_clusters=n_clusters, n_init=10).fit(all_hidden_states)
     state_assignments = kmeans.labels_
     
-    # Step C: DFA Reconstruction
+    # Step C: State Elaboration Logic
+    state_info = "### 🧠 State Interpretation & Dataset Mapping\n"
+    cluster_texts = collections.defaultdict(list)
+    for idx, cluster_id in enumerate(state_assignments):
+        cluster_texts[cluster_id].append(input_snippets[idx])
+        
+    for cluster_id in range(n_clusters):
+        snippets = cluster_texts[cluster_id]
+        # Identify common tokens/attributes that represent this state
+        summary = " | ".join([s[:40] + "..." for s in snippets[:2]])
+        state_info += f"**State S{cluster_id}**: Representing context such as: *{summary}*\n\n"
+
+    # Step D: DFA Reconstruction
     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)
 
     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', arrowsize=20)
-    plt.savefig("dfa_output.png")
-    plt.close() # Clean up memory
+    pos = nx.kamada_kawai_layout(G)
+    nx.draw(G, pos, with_labels=True, node_color='#FF8C00', node_size=3500, font_weight='bold', font_size=12, arrowsize=20)
+    plt.savefig("dfa_output.png", transparent=True)
+    plt.close()
+    
+    analysis_brief = f"Model '{model_name}' identified {n_clusters} distinct equivalence classes in the '{dataset_key}' dataset."
     
-    return "dfa_output.png", f"Model '{model_name}' reduced this dataset into {n_clusters} distinct internal states."
+    return "dfa_output.png", analysis_brief, state_info
 
-# 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.")
+# 2. Gradio UI with Elaboration
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🌐 The Universal Newtonian Probe")
+    gr.Markdown("Extracting the hidden Deterministic Finite Automaton (DFA) from any model and dataset.")
     
     with gr.Row():
         m_drop = gr.Dropdown(choices=MODELS, label="Select Model", value="gpt2")
         d_drop = gr.Dropdown(choices=list(DATASET_CONFIGS.keys()), label="Select Dataset", value="wikitext (v2-raw)")
     
-    btn = gr.Button("Analyze Coherence")
+    btn = gr.Button("Analyze Coherence", variant="primary")
     
     with gr.Row():
-        out_img = gr.Image(label="Extracted DFA")
-        out_txt = gr.Textbox(label="Analysis Result")
+        out_img = gr.Image(label="Extracted DFA (World Map)")
+        with gr.Column():
+            out_txt = gr.Textbox(label="Analysis Status")
+            out_elaboration = gr.Markdown() # Markdown for better readability of interpretation
     
-    btn.click(analyze_world_model, inputs=[m_drop, d_drop], outputs=[out_img, out_txt])
+    btn.click(analyze_world_model, inputs=[m_drop, d_drop], outputs=[out_img, out_txt, out_elaboration])
 
 demo.launch()