Update model_slm.py

model_slm.py

@@ -54,7 +54,7 @@ class MixtureOfRecursionsConfig(PretrainedConfig):
         self.max_position_embeddings = max_position_embeddings or max_seq_len
 
 # ============================================================================
-# EMBEDDINGS MODULE (merged from embeddings.py)
+# EMBEDDINGS MODULE
 # ============================================================================
 
 DEFAULT_BASE = 10000.0
@@ -400,66 +400,113 @@ class RecursiveTransformerLayer(nn.Module):
             active_batches &= (steps > step)
         return x, computation_loss
 
-class MixtureOfRecursions(nn.Module):
+# ============================================================================
+# PRETRAINED MODEL WRAPPER
+# ============================================================================
+
+class MixtureOfRecursionsPreTrainedModel(PreTrainedModel):
+    """PreTrainedModel wrapper for MixtureOfRecursions."""
+    
+    config_class = MixtureOfRecursionsConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    
+    def _init_weights(self, module):
+        """Initialize weights."""
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.d_model ** -0.5)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.d_model ** -0.5)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+class MixtureOfRecursions(MixtureOfRecursionsPreTrainedModel):
     """Transformer model with mixture of recursive layers for technical content."""
     
-    def __init__(
-        self,
-        vocab_size: int,
-        d_model: int = DEFAULT_D_MODEL,
-        n_layers: int = DEFAULT_N_LAYERS,
-        n_heads: int = DEFAULT_N_HEADS,
-        max_steps: int = DEFAULT_MAX_STEPS,
-        dim_feedforward: int = DEFAULT_DIM_FEEDFORWARD,
-        dropout: float = DEFAULT_DROPOUT,
-        max_seq_len: int = DEFAULT_MAX_SEQ_LEN,
-        router_type: str = DEFAULT_ROUTER_TYPE,
-        padding_idx: int = DEFAULT_PADDING_IDX,
-        pos_encoding: str = "learned"
-    ):
-        super().__init__()
-        self.d_model = d_model
-        self.vocab_size = vocab_size
-        self.padding_idx = padding_idx
+    def __init__(self, config: MixtureOfRecursionsConfig):
+        super().__init__(config)
+        self.config = config
+        self.d_model = config.d_model
+        self.vocab_size = config.vocab_size
+        self.padding_idx = config.padding_idx
+        
         self.embeddings = TechEmbeddingLayer(
-            vocab_size=vocab_size,
-            d_model=d_model,
-            max_seq_len=max_seq_len,
-            dropout=dropout,
-            padding_idx=padding_idx,
-            pos_encoding=pos_encoding
+            vocab_size=config.vocab_size,
+            d_model=config.d_model,
+            max_seq_len=config.max_seq_len,
+            dropout=config.dropout,
+            padding_idx=config.padding_idx,
+            pos_encoding=config.pos_encoding
         )
+        
         self.layers = nn.ModuleList([
             RecursiveTransformerLayer(
-                d_model=d_model,
-                n_heads=n_heads,
-                dim_feedforward=dim_feedforward,
-                max_steps=max_steps,
-                dropout=dropout,
-                router_type=router_type
-            ) for _ in range(n_layers)
+                d_model=config.d_model,
+                n_heads=config.n_heads,
+                dim_feedforward=config.dim_feedforward,
+                max_steps=config.max_steps,
+                dropout=config.dropout,
+                router_type=config.router_type
+            ) for _ in range(config.n_layers)
         ])
-        self.final_norm = nn.LayerNorm(d_model)
-        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
-        self._init_weights()
-
-    def _init_weights(self) -> None:
-        nn.init.xavier_uniform_(self.lm_head.weight)
+        
+        self.final_norm = nn.LayerNorm(config.d_model)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        
+        # Initialize weights
+        self.post_init()
 
-    def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: bool = True
+    ):
         batch_size, seq_len = input_ids.shape
+        
+        # Create masks
         padding_mask = create_padding_mask(input_ids, self.padding_idx) if attention_mask is None else (attention_mask == 0)
         causal_mask = create_causal_mask(seq_len, input_ids.device)
         combined_mask = padding_mask.unsqueeze(1).expand(batch_size, seq_len, seq_len) | causal_mask.unsqueeze(0)
+        
+        # Forward pass
         x = self.embeddings(input_ids)
         pos_encoding = self.embeddings.get_positional_encoding()
+        
         total_computation_loss = torch.tensor(0.0, device=x.device)
         for layer in self.layers:
             x, comp_loss = layer(x, combined_mask, pos_encoding)
             total_computation_loss += comp_loss
+        
         x = self.final_norm(x)
         logits = self.lm_head(x)
-        return logits, total_computation_loss
+        
+        loss = None
+        if labels is not None:
+            # Shift logits and labels for language modeling
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(shift_logits.view(-1, self.vocab_size), shift_labels.view(-1))
+            loss += 0.01 * total_computation_loss  # Add computation loss
+        
+        if not return_dict:
+            output = (logits,)
+            return ((loss,) + output) if loss is not None else output
+        
+        from transformers.modeling_outputs import CausalLMOutput
+        return CausalLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=None,
+            attentions=None,
+        )
 
     def generate_step(
         self,
@@ -470,11 +517,14 @@ class MixtureOfRecursions(nn.Module):
     ) -> torch.Tensor:
         self.eval()
         with torch.no_grad():
-            logits, _ = self.forward(input_ids)
+            outputs = self.forward(input_ids, return_dict=True)
+            logits = outputs.logits
             last_logits = logits[:, -1, :] / temperature
+            
             if top_k is not None:
                 indices_to_remove = last_logits < torch.topk(last_logits, top_k)[0][..., -1, None]
                 last_logits = last_logits.masked_fill(indices_to_remove, float('-inf'))
+            
             if top_p is not None:
                 sorted_logits, sorted_indices = torch.sort(last_logits, descending=True)
                 cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
@@ -483,83 +533,12 @@ class MixtureOfRecursions(nn.Module):
                 sorted_indices_to_remove[..., 0] = False
                 indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
                 last_logits = last_logits.masked_fill(indices_to_remove, float('-inf'))
+            
             probs = F.softmax(last_logits, dim=-1)
             return torch.multinomial(probs, num_samples=1)
 
-class TextGenerator:
-    """Text generation utility for the MixtureOfRecursions model."""
-    
-    def __init__(self, model: nn.Module, tokenizer: 'Tokenizer', max_length: int = DEFAULT_MAX_SEQ_LEN, device: Optional[torch.device] = None):
-        self.model = model
-        self.tokenizer = tokenizer
-        self.max_length = max_length
-        self.device = device if device else next(model.parameters()).device
-        self.model.to(self.device)
-        self.eos_token_id = tokenizer.vocab.get('<|endoftext|>', -1)
-        self.assistant_token_id = tokenizer.vocab.get('<|assistant|>', -1)
-
-    def generate(
-        self,
-        prompt: str,
-        method: str = "nucleus",
-        temperature: float = 1.0,
-        top_k: Optional[int] = 50,
-        top_p: Optional[float] = 0.9,
-        max_new_tokens: Optional[int] = None
-    ) -> str:
-        max_new_tokens = max_new_tokens or self.max_length
-        input_text = f"<|user|> {prompt}"
-        input_ids = self.tokenizer.encode_ids(input_text, add_special_tokens=True)
-        input_tensor = torch.tensor([input_ids], device=self.device)
-        self.model.eval()
-        generated_ids = []
-        with torch.no_grad():
-            for _ in range(max_new_tokens):
-                if input_tensor.size(1) > self.max_length:
-                    input_tensor = input_tensor[:, -self.max_length:]
-                if method == "greedy":
-                    next_token = self._greedy_generate(input_tensor)
-                elif method == "sample":
-                    next_token = self._sample_generate(input_tensor, temperature)
-                elif method == "top_k":
-                    next_token = self._top_k_generate(input_tensor, temperature, top_k)
-                elif method == "nucleus" or method == "top_p":
-                    next_token = self._nucleus_generate(input_tensor, temperature, top_p)
-                else:
-                    raise ValueError(f"Unknown generation method: {method}")
-                next_token_id = next_token.item()
-                generated_ids.append(next_token_id)
-                input_tensor = torch.cat([input_tensor, next_token.unsqueeze(0)], dim=1)
-                if next_token_id == self.eos_token_id or (self.assistant_token_id != -1 and next_token_id == self.assistant_token_id):
-                    break
-        full_ids = input_ids + generated_ids
-        full_text = self.tokenizer.decode_ids(full_ids, skip_special_tokens=False)
-        if "<|assistant|>" in full_text:
-            response = full_text.split("<|assistant|>")[-1].split("<|endoftext|>")[0].strip()
-        else:
-            response = full_text.split("<|endoftext|>")[0].strip()
-        return response if response else "No response generated."
-
-    def _greedy_generate(self, input_tensor: torch.Tensor) -> torch.Tensor:
-        logits, _ = self.model(input_tensor)
-        return torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
-
-    def _sample_generate(self, input_tensor: torch.Tensor, temperature: float) -> torch.Tensor:
-        logits, _ = self.model(input_tensor)
-        logits = logits[:, -1, :] / temperature
-        probs = F.softmax(logits, dim=-1)
-        return torch.multinomial(probs, num_samples=1)
-
-    def _top_k_generate(self, input_tensor: torch.Tensor, temperature: float, top_k: int) -> torch.Tensor:
-        logits, _ = self.model(input_tensor)
-        logits = logits[:, -1, :] / temperature
-        top_k_logits, top_k_indices = torch.topk(logits, top_k)
-        probs = F.softmax(top_k_logits, dim=-1)
-        next_token_idx = torch.multinomial(probs, num_samples=1)
-        return top_k_indices.gather(-1, next_token_idx)
-
-    def _nucleus_generate(self, input_tensor: torch.Tensor, temperature: float, top_p: float) -> torch.Tensor:
-        return self.model.generate_step(input_tensor, temperature, top_p=top_p)
+# Register the model for auto class
+MixtureOfRecursions.register_for_auto_class("AutoModelForCausalLM")
 
 def count_parameters(model: nn.Module) -> Tuple[int, int]:
     total_params = sum(p.numel() for p in model.parameters())
@@ -569,7 +548,7 @@ def count_parameters(model: nn.Module) -> Tuple[int, int]:
 def main():
     """Test the MixtureOfRecursions model and its components."""
     print("Initializing MixtureOfRecursions model...")
-    model = MixtureOfRecursions(
+    config = MixtureOfRecursionsConfig(
         vocab_size=DEFAULT_VOCAB_SIZE,
         d_model=DEFAULT_D_MODEL,
         n_layers=DEFAULT_N_LAYERS,
@@ -579,6 +558,7 @@ def main():
         dropout=DEFAULT_DROPOUT,
         router_type=DEFAULT_ROUTER_TYPE
     )
+    model = MixtureOfRecursions(config)
     
     total_params, trainable_params = count_parameters(model)
     print(f"Total parameters: {total_params:,}")
@@ -590,13 +570,13 @@ def main():
     attention_mask = torch.ones_like(input_ids)
     attention_mask[:, -10:] = 0
     
-    logits, comp_loss = model(input_ids, attention_mask)
+    outputs = model(input_ids, attention_mask, return_dict=True)
+    logits = outputs.logits
     
     assert logits.shape == (batch_size, seq_len, DEFAULT_VOCAB_SIZE), f"Unexpected logits shape: {logits.shape}"
     print(f"Input shape: {input_ids.shape}")
     print(f"Output logits shape: {logits.shape}")
     print(f"Expected logits shape: ({batch_size}, {seq_len}, {DEFAULT_VOCAB_SIZE})")
-    print(f"Computation loss: {comp_loss:.4f}")
     
     print("\nTesting generation step...")
     next_token = model.generate_step(input_ids[:1], temperature=0.8, top_p=0.9)