Update memory.py

memory.py

@@ -17,11 +17,13 @@ class CognitiveMemory(nn.Module):
         self.consolidation_threshold = 0.7
         
         # Memory projection layers
-        self.key_proj = nn.Linear(context_size, 64)
-        self.value_proj = nn.Linear(context_size, 64)
+        self.key_proj = nn.Linear(1, 64)  # Changed from context_size to 1
+        self.value_proj = nn.Linear(1, 64)  # Changed from context_size to 1
         
     def add_memory(self, context: torch.Tensor, activation: float):
         """Store new memory with adaptive importance"""
+        # Ensure context is 1D tensor with single value
+        context = context.reshape(-1)
         importance = torch.sigmoid(torch.tensor(activation * 0.5 + 0.2))
         self.memory_queue.append({
             'context': context.detach(),
@@ -46,9 +48,14 @@ class CognitiveMemory(nn.Module):
         if not self.memory_queue:
             return torch.zeros_like(query)
             
-        keys = torch.stack([self.key_proj(m['context']) for m in self.memory_queue])
-        values = torch.stack([self.value_proj(m['context']) for m in self.memory_queue])
+        # Ensure query is 1D tensor with single value
+        query = query.reshape(1, 1)
+        memories = torch.stack([m['context'].reshape(1, 1) for m in self.memory_queue])
+        
+        keys = self.key_proj(memories)
+        values = self.value_proj(memories)
         query_proj = self.key_proj(query)
         
-        scores = F.softmax(keys @ query_proj.t(), dim=0)
-        return (scores * values).sum(dim=0)
+        scores = F.softmax(torch.matmul(keys, query_proj.transpose(0, 1)), dim=0)
+        retrieved = torch.matmul(scores.transpose(0, 1), values)
+        return retrieved.squeeze(0)