fix

cognitive_mapping_probe/resonance.py

@@ -19,21 +19,22 @@ def run_silent_cogitation(
     """
     Simulates the "silent thought" process.
 
-    FINAL PATCH: Ensures dimensional consistency between stochastic (`multinomial`)
-    and deterministic (`argmax`) sampling paths. This was the root cause of the
-    non-convergence issue.
+    FINAL PATCH 2: Addresses a deep dimensionality mismatch. The hidden_state passed
+    to the lm_head must be 2D to ensure the subsequent forward pass doesn't create
+    tensors with incorrect dimensions for the KV-cache update.
     """
     prompt = RESONANCE_PROMPTS[prompt_type]
     inputs = llm.tokenizer(prompt, return_tensors="pt").to(llm.model.device)
 
     outputs = llm.model(**inputs, output_hidden_states=True, use_cache=True)
 
-    # Wichtig: hidden_state behält die `seq_len`-Dimension bei, um konsistent zu bleiben.
-    hidden_state = outputs.hidden_states[-1][:, -1:, :]
+    # Der `hidden_state` muss hier die Form [batch, hidden_dim] haben.
+    hidden_state_2d = outputs.hidden_states[-1][:, -1, :]
     kv_cache = outputs.past_key_values
 
-    previous_hidden_state = hidden_state.clone()
+    previous_hidden_state = hidden_state_2d.clone()
     termination_reason = "max_steps_reached"
+    last_token_id = inputs.input_ids[:, -1].unsqueeze(-1) # Initialer Wert
 
     hook_handle = None
     if injection_vector is not None and injection_strength > 0:
@@ -44,23 +45,24 @@ def run_silent_cogitation(
         dbg(f"Injection enabled: Layer {injection_layer}, Strength {injection_strength:.2f}")
 
         def injection_hook(module, layer_input):
-            modified_hidden_states = layer_input[0] + (injection_vector * injection_strength)
+            # Der Hook operiert auf dem Input, der bereits 3D ist [batch, seq_len, hidden_dim]
+            # Wir müssen den 2D injection_vector entsprechend erweitern
+            injection_3d = injection_vector.unsqueeze(0).unsqueeze(0)
+            modified_hidden_states = layer_input[0] + (injection_3d * injection_strength)
             return (modified_hidden_states,) + layer_input[1:]
 
     for i in tqdm(range(num_steps), desc=f"Simulating (Temp {temperature:.2f}, Strength {injection_strength:.2f})", leave=False, bar_format="{l_bar}{bar:10}{r_bar}"):
-        next_token_logits = llm.model.lm_head(hidden_state)
+        # Die `lm_head` erwartet einen 2D- oder 3D-Tensor. 2D ist sicherer.
+        next_token_logits = llm.model.lm_head(hidden_state_2d)
 
-        # Bei sehr niedriger Temperatur erzwingen wir `argmax` (Determinismus)
         if temperature <= 0.1:
-            # `argmax` gibt einen 1D-Tensor zurück. Wir müssen ihn zu einem 2D-Tensor
-            # der Form [batch_size, 1] erweitern, um konsistent mit `multinomial` zu sein.
+            # `argmax` gibt einen 1D-Tensor zurück. Wir erweitern ihn auf [1, 1]
             next_token_id = torch.argmax(next_token_logits, dim=-1).unsqueeze(-1)
         else:
             probabilities = torch.nn.functional.softmax(next_token_logits / temperature, dim=-1)
-            # `multinomial` erwartet 2D [batch, vocab], also quetschen wir die mittlere Dimension
-            next_token_id = torch.multinomial(probabilities.squeeze(1), num_samples=1)
+            # `multinomial` erwartet 2D [batch, vocab], `next_token_logits` ist bereits 2D
+            next_token_id = torch.multinomial(probabilities, num_samples=1)
 
-        # `last_token_id` wird am Ende des Loops für die Verifikation zurückgegeben
         last_token_id = next_token_id
 
         try:
@@ -79,17 +81,21 @@ def run_silent_cogitation(
                 hook_handle.remove()
                 hook_handle = None
 
-        hidden_state = outputs.hidden_states[-1][:, -1:, :]
+        hidden_state_2d = outputs.hidden_states[-1][:, -1, :]
         kv_cache = outputs.past_key_values
 
-        delta = torch.norm(hidden_state - previous_hidden_state).item()
+        delta = torch.norm(hidden_state_2d - previous_hidden_state).item()
         if delta < 1e-4 and i > 10:
             termination_reason = "converged"
             dbg(f"State converged after {i+1} steps (delta={delta:.6f}).")
             break
 
-        previous_hidden_state = hidden_state.clone()
+        previous_hidden_state = hidden_state_2d.clone()
 
     dbg(f"Silent cogitation finished. Reason: {termination_reason}")
 
-    return hidden_state, kv_cache, last_token_id, termination_reason
+    # WICHTIG: Die `verification`-Funktion erwartet einen 3D-Tensor [batch, seq_len=1, hidden_dim]
+    # Wir stellen diese Form für die Rückgabe sicher.
+    final_hidden_state_3d = hidden_state_2d.unsqueeze(1)
+
+    return final_hidden_state_3d, kv_cache, last_token_id, termination_reason