Add CognitiveProcessIntegration and InformationIntegration classes; i…

…mplement tests for ARC reasoning

models/consciousness_model.py

@@ -205,3 +205,90 @@ def forward(self, inputs, deterministic=False, initial_state=None):
         memory_output = self.layer_norm(memory_output)
 
         return memory_output, memory_state
+
+class CognitiveProcessIntegration(nn.Module):
+    def __init__(self, hidden_dim: int, num_heads: int, dropout_rate: float):
+        super().__init__()
+        self.hidden_dim = hidden_dim
+        self.num_heads = num_heads
+        self.attention = nn.MultiheadAttention(
+            embed_dim=hidden_dim,
+            num_heads=num_heads,
+            dropout=dropout_rate,
+            batch_first=True
+        )
+        self.input_projection = nn.Linear(hidden_dim, hidden_dim)  # Added input_projection
+        # ...initialize other necessary layers...
+
+    def forward(self, inputs: Dict[str, torch.Tensor], deterministic: bool = True):
+        """Process multiple modalities and generate cross-modal attention maps."""
+        batch_size = next(iter(inputs.values())).size(0)
+        seq_length = next(iter(inputs.values())).size(1)
+        attention_maps = {}
+        processed_states = {}
+
+        # First pass: Project all inputs
+        for modality, tensor in inputs.items():
+            processed = self.input_projection(tensor)  # Use input_projection
+            processed_states[modality] = processed
+
+        # Initialize combined state with zeros matching the maximum sequence length
+        max_seq_length = max(tensor.size(1) for tensor in processed_states.values())
+        combined_state = torch.zeros(
+            batch_size, max_seq_length, self.hidden_dim,
+            device=next(iter(inputs.values())).device
+        )
+
+        # Generate attention maps between all modality pairs
+        for source in inputs.keys():
+            for target in inputs.keys():
+                if source != target:
+                    query = processed_states[target]
+                    key = processed_states[source]
+                    value = processed_states[source]
+
+                    attn_output, attn_weights = self.attention(
+                        query=query,
+                        key=key,
+                        value=value
+                    )
+
+                    # Store attention map
+                    map_key = f"{target}-{source}"
+                    attention_maps[map_key] = attn_weights
+
+                    # Pad attn_output if necessary to match combined_state's sequence length
+                    if attn_output.size(1) < max_seq_length:
+                        pad_size = max_seq_length - attn_output.size(1)
+                        attn_output = torch.nn.functional.pad(attn_output, (0, 0, 0, pad_size))
+                    elif attn_output.size(1) > max_seq_length:
+                        attn_output = attn_output[:, :max_seq_length, :]
+
+                    combined_state = combined_state + attn_output
+
+        # ...existing code...
+        return combined_state, attention_maps
+
+class InformationIntegration(nn.Module):
+    def __init__(self, hidden_dim: int, num_modules: int, dropout_rate: float):
+        super().__init__()
+        # Store modules in a ModuleList
+        self.module_list = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(hidden_dim, hidden_dim),
+                nn.ReLU(),
+                nn.Dropout(dropout_rate)
+            ) for _ in range(num_modules)
+        ])
+        self.phi_layer = nn.Linear(hidden_dim, 1)
+
+    def forward(self, memory_output: torch.Tensor, deterministic: bool = True):
+        integrated_output = memory_output
+        # Iterate through module_list instead of calling modules()
+        for module in self.module_list:
+            integrated_output = module(integrated_output)
+
+        # Compute phi with non-linearity to introduce variability
+        phi = torch.sigmoid(self.phi_layer(integrated_output)).squeeze(-1)
+
+        return integrated_output, phi

tests/benchmarks/test_arc_reasoning.py

@@ -0,0 +1,161 @@
+import torch
+import pytest
+from typing import Dict, Tuple
+
+from models.consciousness_model import ConsciousnessModel
+
+class TestARCReasoning:
+    @pytest.fixture
+    def device(self):
+        return torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    @pytest.fixture
+    def model_config(self):
+        return ConsciousnessModel.create_default_config()
+
+    @pytest.fixture
+    def consciousness_model(self, model_config):
+        return ConsciousnessModel(**model_config)
+
+    def load_arc_sample(self) -> Tuple[Dict[str, torch.Tensor], torch.Tensor]:
+        """Load a sample ARC task for testing."""
+        sample_input = {
+            'visual': torch.tensor([
+                [1, 0, 1],
+                [0, 1, 0],
+                [1, 0, 0]
+            ], dtype=torch.float32).unsqueeze(0).unsqueeze(-1)
+        }
+
+        expected_output = torch.tensor([
+            [1, 0, 1],
+            [0, 1, 0],
+            [1, 0, 1]
+        ], dtype=torch.float32).unsqueeze(0).unsqueeze(-1)
+
+        return sample_input, expected_output
+
+    def _prepare_visual_input(self, visual, batch_size, hidden_dim):
+        """Prepare visual input for the model by flattening and padding."""
+        visual_flat = visual.view(batch_size, -1)  # Flatten to [batch_size, N]
+        return torch.nn.functional.pad(
+            visual_flat,
+            (0, hidden_dim - visual_flat.shape[1])
+        )
+
+    def _get_final_state(self, output):
+        """Extract final state from model output."""
+        if output.dim() == 3:  # If output is [batch, seq, hidden]
+            return output[:, -1, :]  # Take last sequence position
+        return output  # Otherwise return as is
+
+    def test_pattern_recognition(self, device, consciousness_model):
+        inputs, expected = self.load_arc_sample()
+        batch_size = inputs['visual'].shape[0]
+
+        # Project visual input to correct dimensionality
+        visual_input = self._prepare_visual_input(
+            inputs['visual'], 
+            batch_size,
+            consciousness_model.hidden_dim
+        )
+
+        # Initialize model state
+        model_inputs = {
+            'visual': visual_input.to(device),
+            'state': torch.zeros((batch_size, consciousness_model.hidden_dim), device=device)
+        }
+
+        consciousness_model = consciousness_model.to(device)
+        consciousness_model.eval()
+
+        try:
+            with torch.no_grad():
+                output, metrics = consciousness_model(
+                    model_inputs,
+                    deterministic=True,
+                    consciousness_threshold=0.5
+                )
+
+                final_state = self._get_final_state(output)
+
+                # Validate outputs
+                assert final_state.shape == (batch_size, consciousness_model.hidden_dim)
+                assert 'phi' in metrics
+                assert metrics['phi'].shape == (batch_size, 1)
+                assert torch.all(metrics['phi'] >= 0)
+
+                # Validate attention
+                assert 'attention_weights' in metrics
+                assert metrics['attention_weights'].dim() >= 3  # (batch, heads, seq)
+
+                # Validate attention maps
+                assert 'attention_maps' in metrics
+                for attn_map in metrics['attention_maps'].values():
+                    assert torch.allclose(
+                        torch.sum(attn_map, dim=-1),
+                        torch.ones((batch_size, 8, 64), device=device)
+                    )
+
+        except Exception as e:
+            pytest.fail(f"Pattern recognition test failed: {str(e)}")
+
+    def test_abstraction_capability(self, device, consciousness_model):
+        inputs, _ = self.load_arc_sample()
+        batch_size = inputs['visual'].shape[0]
+
+        # Create transformed versions
+        def preprocess_input(x):
+            return self._prepare_visual_input(
+                x, 
+                batch_size,
+                consciousness_model.hidden_dim
+            )
+
+        variations = {
+            'original': preprocess_input(inputs['visual']),
+            'rotated': preprocess_input(torch.rot90(inputs['visual'][:, :, :, 0], k=1).unsqueeze(-1)),
+            'scaled': preprocess_input(inputs['visual'] * 2.0)
+        }
+
+        consciousness_model = consciousness_model.to(device)
+        consciousness_model.eval()
+
+        try:
+            states = {}
+            with torch.no_grad():
+                for name, visual_input in variations.items():
+                    output, metrics = consciousness_model(
+                        {'visual': visual_input.to(device),
+                         'state': torch.zeros((batch_size, consciousness_model.hidden_dim), device=device)},
+                        deterministic=True
+                    )
+                    states[name] = self._get_final_state(output)
+
+            # Test representation similarity
+            def cosine_similarity(x, y):
+                return torch.sum(x * y) / (torch.linalg.norm(x) * torch.linalg.norm(y))
+
+            orig_rot_sim = cosine_similarity(
+                states['original'].flatten(),
+                states['rotated'].flatten()
+            )
+            orig_scaled_sim = cosine_similarity(
+                states['original'].flatten(),
+                states['scaled'].flatten()
+            )
+
+            # Transformed versions should maintain similar representations
+            assert orig_rot_sim > 0.5
+            assert orig_scaled_sim > 0.7
+
+        except Exception as e:
+            pytest.fail(f"Abstraction capability test failed: {str(e)}")
+
+    # Convert remaining test methods similarly...
+    # The pattern is similar - main changes are:
+    # - Replace jnp with torch
+    # - Use .to(device) for tensors
+    # - Use torch.no_grad() instead of JAX's deterministic flag
+    # - Use PyTorch's tensor operations instead of JAX's
+