Upload model_architecture.py

model_architecture.py

@@ -90,527 +90,4 @@ class ShortTermMemory(nn.Module):
             'memory_read': memory_read,
             'read_weights': read_weights,
             'updated_memory': memory
-        }
-
-
-class LongTermMemory(nn.Module):
-    """
-    Long-term memory module for persistent storage of knowledge.
-    Uses a key-value store with retrieval mechanism.
-    
-    Inspired by:
-    - LongMem's adaptive residual side-network
-    - MemoRAG's global memory for overview
-    """
-    
-    def __init__(self, hidden_size: int, num_slots: int = 10000):
-        super().__init__()
-        self.hidden_size = hidden_size
-        self.num_slots = num_slots
-        
-        # Memory storage (key-value pairs)
-        self.memory_keys = nn.Parameter(torch.zeros(num_slots, hidden_size))
-        self.memory_values = nn.Parameter(torch.zeros(num_slots, hidden_size))
-        
-        # Memory metadata (importance, timestamps, etc.)
-        self.memory_importance = nn.Parameter(torch.ones(num_slots))
-        
-        # Query projection
-        self.query_proj = nn.Linear(hidden_size, hidden_size)
-        self.key_proj = nn.Linear(hidden_size, hidden_size)
-        self.value_proj = nn.Linear(hidden_size, hidden_size)
-        
-        # Memory consolidation (for summarizing old memories)
-        self.consolidation_layer = nn.Sequential(
-            nn.Linear(hidden_size * 2, hidden_size),
-            nn.ReLU(),
-            nn.Linear(hidden_size, hidden_size)
-        )
-        
-        # Initialize
-        nn.init.normal_(self.memory_keys, mean=0, std=0.02)
-        nn.init.normal_(self.memory_values, mean=0, std=0.02)
-        nn.init.constant_(self.memory_importance, 1.0)
-    
-    def forward(self, hidden_states: torch.Tensor,
-                store_new: bool = False,
-                new_content: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
-        """
-        Retrieve from or store to long-term memory.
-        
-        Args:
-            hidden_states: Current hidden states [batch_size, seq_len, hidden_size]
-            store_new: Whether to store new content
-            new_content: New content to store [batch_size, content_len, hidden_size]
-            
-        Returns:
-            Dictionary with memory outputs
-        """
-        batch_size, seq_len, _ = hidden_states.shape
-        
-        # Project queries
-        queries = self.query_proj(hidden_states)  # [batch, seq, hidden]
-        
-        # Project memory keys
-        keys = self.key_proj(self.memory_keys)  # [num_slots, hidden]
-        
-        # Compute retrieval scores
-        scores = torch.matmul(queries, keys.transpose(-1, -2))  # [batch, seq, num_slots]
-        scores = scores / math.sqrt(self.hidden_size)
-        
-        # Apply importance weighting
-        importance_weights = torch.sigmoid(self.memory_importance)  # [num_slots]
-        scores = scores * importance_weights.unsqueeze(0).unsqueeze(0)  # [batch, seq, num_slots]
-        
-        # Get top-k memories
-        top_k = min(10, self.num_slots)
-        top_scores, top_indices = torch.topk(scores, k=top_k, dim=-1)  # [batch, seq, top_k]
-        
-        # Softmax over top-k
-        retrieval_weights = F.softmax(top_scores, dim=-1)  # [batch, seq, top_k]
-        
-        # Gather retrieved memories
-        retrieved_values = self.memory_values[top_indices]  # [batch, seq, top_k, hidden]
-        
-        # Weighted sum
-        memory_output = torch.sum(retrieved_values * retrieval_weights.unsqueeze(-1), dim=-2)  # [batch, seq, hidden]
-        
-        # Project memory output
-        memory_output = self.value_proj(memory_output)
-        
-        # Store new memories if requested
-        if store_new and new_content is not None:
-            self._store_new_memories(new_content)
-        
-        return {
-            'output': memory_output,
-            'retrieval_weights': retrieval_weights,
-            'top_indices': top_indices,
-            'retrieval_scores': top_scores
-        }
-    
-    def _store_new_memories(self, content: torch.Tensor):
-        """
-        Store new content in long-term memory.
-        Uses a simple FIFO with consolidation strategy.
-        
-        Args:
-            content: Content to store [batch_size, content_len, hidden_size]
-        """
-        with torch.no_grad():
-            # Average content across batch and sequence
-            content_repr = content.mean(dim=(0, 1))  # [hidden_size]
-            
-            # Find least important slot
-            _, least_important_idx = torch.min(self.memory_importance, dim=0)
-            
-            # Update memory (with small learning rate)
-            lr = 0.01
-            self.memory_keys.data[least_important_idx] = \
-                (1 - lr) * self.memory_keys.data[least_important_idx] + lr * content_repr
-            self.memory_values.data[least_important_idx] = \
-                (1 - lr) * self.memory_values.data[least_important_idx] + lr * content_repr
-            
-            # Reset importance (new memories start neutral)
-            self.memory_importance.data[least_important_idx] = 1.0
-    
-    def consolidate_memories(self, threshold: float = 0.1):
-        """
-        Consolidate similar memories to prevent redundancy.
-        
-        Args:
-            threshold: Similarity threshold for consolidation
-        """
-        with torch.no_grad():
-            # Compute pairwise similarity
-            keys_norm = F.normalize(self.memory_keys, dim=-1)
-            similarity = torch.matmul(keys_norm, keys_norm.transpose(-1, -2))
-            
-            # Find highly similar pairs
-            mask = torch.triu(similarity > threshold, diagonal=1)
-            
-            if mask.any():
-                # Consolidate: merge similar memories
-                indices = torch.where(mask)
-                
-                for i, j in zip(indices[0], indices[1]):
-                    # Merge j into i (weighted by importance)
-                    imp_i = torch.sigmoid(self.memory_importance[i])
-                    imp_j = torch.sigmoid(self.memory_importance[j])
-                    total_imp = imp_i + imp_j
-                    
-                    if total_imp > 0:
-                        self.memory_keys.data[i] = \
-                            (imp_i * self.memory_keys.data[i] + imp_j * self.memory_keys.data[j]) / total_imp
-                        self.memory_values.data[i] = \
-                            (imp_i * self.memory_values.data[i] + imp_j * self.memory_values.data[j]) / total_imp
-                        self.memory_importance.data[i] = torch.log(torch.exp(self.memory_importance.data[i]) + 
-                                                                 torch.exp(self.memory_importance.data[j]))
-                        
-                        # Zero out merged memory
-                        self.memory_keys.data[j] = 0
-                        self.memory_values.data[j] = 0
-                        self.memory_importance.data[j] = -10  # Mark as inactive
-
-
-class MemoryAugmentedLayer(nn.Module):
-    """
-    Transformer layer augmented with both short-term and long-term memory.
-    """
-    
-    def __init__(self, hidden_size: int, num_heads: int, 
-                 short_term_memory_size: int = 512,
-                 long_term_memory_slots: int = 10000):
-        super().__init__()
-        self.hidden_size = hidden_size
-        self.num_heads = num_heads
-        
-        # Self-attention
-        self.attention = nn.MultiheadAttention(hidden_size, num_heads, batch_first=True)
-        self.attention_norm = nn.LayerNorm(hidden_size)
-        
-        # Feed-forward
-        self.ffn = nn.Sequential(
-            nn.Linear(hidden_size, hidden_size * 4),
-            nn.GELU(),
-            nn.Linear(hidden_size * 4, hidden_size)
-        )
-        self.ffn_norm = nn.LayerNorm(hidden_size)
-        
-        # Memory modules
-        self.short_term_memory = ShortTermMemory(hidden_size, short_term_memory_size)
-        self.long_term_memory = LongTermMemory(hidden_size, long_term_memory_slots)
-        
-        # Memory gating
-        self.memory_gate = nn.Linear(hidden_size * 2, hidden_size)
-    
-    def forward(self, hidden_states: torch.Tensor,
-                attention_mask: Optional[torch.Tensor] = None,
-                use_long_term_memory: bool = True,
-                store_long_term: bool = False) -> Dict[str, torch.Tensor]:
-        """
-        Forward pass through memory-augmented layer.
-        
-        Args:
-            hidden_states: Input hidden states [batch_size, seq_len, hidden_size]
-            attention_mask: Optional attention mask
-            use_long_term_memory: Whether to use long-term memory
-            store_long_term: Whether to store in long-term memory
-            
-        Returns:
-            Dictionary with layer outputs
-        """
-        residual = hidden_states
-        
-        # Self-attention
-        attn_output, attn_weights = self.attention(
-            self.attention_norm(hidden_states),
-            self.attention_norm(hidden_states),
-            self.attention_norm(hidden_states),
-            key_padding_mask=attention_mask
-        )
-        hidden_states = residual + attn_output
-        
-        # Short-term memory
-        memory_output = self.short_term_memory(hidden_states, attention_mask)
-        hidden_states = memory_output['output']
-        
-        # Long-term memory (optional)
-        long_term_output = None
-        if use_long_term_memory:
-            long_term_output = self.long_term_memory(
-                hidden_states,
-                store_new=store_long_term,
-                new_content=residual  # Store original input
-            )
-            # Gate between attention and memory
-            memory_gate_input = torch.cat([hidden_states, long_term_output['output']], dim=-1)
-            gate = torch.sigmoid(self.memory_gate(memory_gate_input))
-            hidden_states = gate * hidden_states + (1 - gate) * long_term_output['output']
-        
-        # Feed-forward
-        ffn_output = self.ffn(self.ffn_norm(hidden_states))
-        hidden_states = hidden_states + ffn_output
-        
-        return {
-            'hidden_states': hidden_states,
-            'attention_weights': attn_weights,
-            'short_term_output': memory_output,
-            'long_term_output': long_term_output
-        }
-
-
-class LMCODE(nn.Module):
-    """
-    Language Model with Memory CODE (LMCODE).
-    
-    A transformer-based language model augmented with:
-    1. Short-term memory: Working memory for immediate context
-    2. Long-term memory: Persistent storage for knowledge and experiences
-    
-    Architecture:
-    - Embedding layer
-    - Multiple memory-augmented transformer layers
-    - Language model head
-    
-    Memory Flow:
-    - Short-term: Updated every forward pass, decays over time
-    - Long-term: Consolidated periodically, stores important experiences
-    """
-    
-    def __init__(self, config: Optional['LMCODEConfig'] = None):
-        super().__init__()
-        
-        if config is None:
-            config = LMCODEConfig()
-        
-        self.config = config
-        
-        # Embeddings
-        self.token_embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.position_embeddings = nn.Embedding(2048, config.hidden_size)
-        
-        # Memory-augmented layers
-        self.layers = nn.ModuleList([
-            MemoryAugmentedLayer(
-                hidden_size=config.hidden_size,
-                num_heads=config.num_heads,
-                short_term_memory_size=config.short_term_memory_size,
-                long_term_memory_slots=config.long_term_memory_slots
-            )
-            for _ in range(config.num_layers)
-        ])
-        
-        # Output normalization
-        self.final_norm = nn.LayerNorm(config.hidden_size)
-        
-        # Language model head
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        
-        # Initialize weights
-        self.apply(self._init_weights)
-    
-    def _init_weights(self, module):
-        """Initialize model weights."""
-        if isinstance(module, nn.Linear):
-            nn.init.normal_(module.weight, mean=0, std=0.02)
-            if module.bias is not None:
-                nn.init.zeros_(module.bias)
-        elif isinstance(module, nn.Embedding):
-            nn.init.normal_(module.weight, mean=0, std=0.02)
-        elif isinstance(module, nn.LayerNorm):
-            nn.init.zeros_(module.bias)
-            nn.init.ones_(module.weight)
-    
-    def forward(self, input_ids: torch.Tensor,
-                attention_mask: Optional[torch.Tensor] = None,
-                use_long_term_memory: bool = True,
-                store_long_term: bool = False,
-                labels: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
-        """
-        Forward pass through LMCODE model.
-        
-        Args:
-            input_ids: Input token IDs [batch_size, seq_len]
-            attention_mask: Optional attention mask [batch_size, seq_len]
-            use_long_term_memory: Whether to use long-term memory
-            store_long_term: Whether to store in long-term memory
-            labels: Optional labels for computing loss
-            
-        Returns:
-            Dictionary with model outputs
-        """
-        batch_size, seq_len = input_ids.shape
-        
-        # Create position IDs
-        position_ids = torch.arange(seq_len, device=input_ids.device).unsqueeze(0).expand(batch_size, -1)
-        
-        # Embeddings
-        token_embeds = self.token_embeddings(input_ids)
-        position_embeds = self.position_embeddings(position_ids)
-        hidden_states = token_embeds + position_embeds
-        
-        # Apply dropout
-        hidden_states = F.dropout(hidden_states, p=0.1, training=self.training)
-        
-        # Track intermediate outputs
-        all_hidden_states = []
-        all_attention_weights = []
-        all_short_term_outputs = []
-        all_long_term_outputs = []
-        
-        # Pass through layers
-        for layer in self.layers:
-            layer_output = layer(
-                hidden_states,
-                attention_mask=attention_mask,
-                use_long_term_memory=use_long_term_memory,
-                store_long_term=store_long_term
-            )
-            
-            hidden_states = layer_output['hidden_states']
-            all_hidden_states.append(hidden_states)
-            all_attention_weights.append(layer_output['attention_weights'])
-            all_short_term_outputs.append(layer_output['short_term_output'])
-            all_long_term_outputs.append(layer_output['long_term_output'])
-        
-        # Final normalization
-        hidden_states = self.final_norm(hidden_states)
-        
-        # Language model logits
-        logits = self.lm_head(hidden_states)
-        
-        # Compute loss if labels provided
-        loss = None
-        if labels is not None:
-            loss = F.cross_entropy(
-                logits.view(-1, self.config.vocab_size),
-                labels.view(-1),
-                ignore_index=-100
-            )
-        
-        return {
-            'logits': logits,
-            'loss': loss,
-            'hidden_states': all_hidden_states,
-            'attention_weights': all_attention_weights,
-            'short_term_outputs': all_short_term_outputs,
-            'long_term_outputs': all_long_term_outputs
-        }
-    
-    def generate(self, input_ids: torch.Tensor,
-                 max_length: int = 100,
-                 temperature: float = 1.0,
-                 top_k: int = 50,
-                 top_p: float = 0.9,
-                 use_long_term_memory: bool = True,
-                 store_long_term: bool = True) -> torch.Tensor:
-        """
-        Generate text autoregressively.
-        
-        Args:
-            input_ids: Initial token IDs [batch_size, seq_len]
-            max_length: Maximum generation length
-            temperature: Sampling temperature
-            top_k: Top-k sampling
-            top_p: Nucleus sampling
-            use_long_term_memory: Whether to use long-term memory during generation
-            store_long_term: Whether to store generated text in long-term memory
-            
-        Returns:
-            Generated token IDs [batch_size, total_length]
-        """
-        batch_size = input_ids.shape[0]
-        
-        for _ in range(max_length):
-            # Forward pass
-            with torch.no_grad():
-                outputs = self(
-                    input_ids,
-                    use_long_term_memory=use_long_term_memory,
-                    store_long_term=store_long_term
-                )
-                
-                logits = outputs['logits'][:, -1, :] / temperature
-                
-                # Top-k filtering
-                if top_k > 0:
-                    indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
-                    logits[indices_to_remove] = -float('Inf')
-                
-                # Top-p (nucleus) filtering
-                if top_p < 1.0:
-                    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
-                    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
-                    
-                    sorted_indices_to_remove = cumulative_probs > top_p
-                    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
-                    sorted_indices_to_remove[..., 0] = 0
-                    
-                    indices_to_remove = sorted_indices_to_remove.scatter(
-                        1, sorted_indices, sorted_indices_to_remove
-                    )
-                    logits[indices_to_remove] = -float('Inf')
-                
-                # Sample next token
-                probs = F.softmax(logits, dim=-1)
-                next_token = torch.multinomial(probs, num_samples=1)
-                
-                # Append to sequence
-                input_ids = torch.cat([input_ids, next_token], dim=-1)
-        
-        return input_ids
-    
-    def store_experience(self, text: str, metadata: Optional[Dict] = None):
-        """
-        Store a text experience in long-term memory.
-        
-        Args:
-            text: Text to store
-            metadata: Optional metadata
-        """
-        # Tokenize text (simplified - in practice, use proper tokenizer)
-        # For demonstration, we'll create dummy embeddings
-        # In production, use: tokens = tokenizer.encode(text)
-        
-        # Create a representation (in practice, pass through model)
-        dummy_input = torch.zeros(1, 10, self.config.hidden_size)
-        
-        with torch.no_grad():
-            # Get representation from model
-            representation = self.forward(dummy_input, use_long_term_memory=False)
-            # Store in long-term memory
-            self.layers[0].long_term_memory._store_new_memories(representation['hidden_states'][-1])
-    
-    def query_memory(self, query: str, top_k: int = 5) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Query long-term memory for relevant information.
-        
-        Args:
-            query: Query text
-            top_k: Number of results to retrieve
-            
-        Returns:
-            retrieved_memories: Retrieved memory representations
-            indices: Indices of retrieved memories
-        """
-        # Create query representation
-        dummy_input = torch.zeros(1, 10, self.config.hidden_size)
-        
-        with torch.no_grad():
-            query_representation = self.forward(dummy_input, use_long_term_memory=False)
-            query_repr = query_representation['hidden_states'][-1].mean(dim=1)
-            
-            # Retrieve from long-term memory
-            retrieved, indices = self.layers[0].long_term_memory.retrieve(query_repr, top_k=top_k)
-        
-        return retrieved, indices
-
-
-class LMCODEConfig:
-    """Configuration for LMCODE model."""
-    
-    def __init__(self, vocab_size: int = 50257, hidden_size: int = 512, 
-                 num_layers: int = 6, num_heads: int = 8,
-                 short_term_memory_size: int = 512,
-                 long_term_memory_slots: int = 10000):
-        self.vocab_size = vocab_size
-        self.hidden_size = hidden_size
-        self.num_layers = num_layers
-        self.num_heads = num_heads
-        self.short_term_memory_size = short_term_memory_size
-        self.long_term_memory_slots = long_term_memory_slots
-    
-    def to_dict(self) -> Dict:
-        return {
-            'vocab_size': self.vocab_size,
-            'hidden_size': self.hidden_size,
-            'num_layers': self.num_layers,
-            'num_heads': self.num_heads,
-            'short_term_memory_size': self.short_term_memory_size,
-            'long_term_memory_slots': self.long_term_memory_slots
-        }
-    
-    @classmethod
-    def from_dict(cls, config_dict: Dict) -> 'LMCODEConfig':
-        return cls(**config_dict)
+        }