Merge embeddings and add Transformers support

model_slm.py

@@ -2,11 +2,10 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import math
-from typing import Optional, Tuple, Union
-from embeddings import TechEmbeddingLayer, create_padding_mask, create_causal_mask
+from typing import Optional, Tuple, Union, List
 
 # ============================================================================
-# TRANSFORMERS COMPATIBILITY - ADD THIS SECTION
+# TRANSFORMERS COMPATIBILITY
 # ============================================================================
 from transformers import PretrainedConfig
 from transformers.modeling_utils import PreTrainedModel
@@ -29,7 +28,6 @@ class MixtureOfRecursionsConfig(PretrainedConfig):
         router_type="adaptive",
         padding_idx=0,
         pos_encoding="learned",
-        # Transformers standard names (for compatibility)
         hidden_size=None,
         num_hidden_layers=None,
         num_attention_heads=None,
@@ -38,8 +36,6 @@ class MixtureOfRecursionsConfig(PretrainedConfig):
         **kwargs
     ):
         super().__init__(**kwargs)
-        
-        # Your model's parameters
         self.vocab_size = vocab_size
         self.d_model = d_model
         self.n_layers = n_layers
@@ -51,8 +47,6 @@ class MixtureOfRecursionsConfig(PretrainedConfig):
         self.router_type = router_type
         self.padding_idx = padding_idx
         self.pos_encoding = pos_encoding
-        
-        # Transformers standard aliases (for compatibility)
         self.hidden_size = hidden_size or d_model
         self.num_hidden_layers = num_hidden_layers or n_layers
         self.num_attention_heads = num_attention_heads or n_heads
@@ -60,10 +54,155 @@ class MixtureOfRecursionsConfig(PretrainedConfig):
         self.max_position_embeddings = max_position_embeddings or max_seq_len
 
 # ============================================================================
-# END TRANSFORMERS COMPATIBILITY SECTION
+# EMBEDDINGS MODULE (merged from embeddings.py)
+# ============================================================================
+
+DEFAULT_BASE = 10000.0
+DEFAULT_CUTOFFS = [2000, 10000]
+DEFAULT_DIV_VAL = 4.0
+
+class PositionalEncoding(nn.Module):
+    """Sinusoidal positional encoding for transformer models."""
+    
+    def __init__(self, d_model: int, max_seq_len: int = 512, dropout: float = 0.1):
+        super().__init__()
+        self.d_model = d_model
+        self.dropout = nn.Dropout(dropout)
+        pe = torch.zeros(max_seq_len, d_model)
+        position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(DEFAULT_BASE) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term[:, :-1] if d_model % 2 == 1 else div_term)
+        self.register_buffer('pe', pe.unsqueeze(0))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, d_model = x.size()
+        if d_model != self.d_model:
+            raise ValueError(f"Input dimension {d_model} does not match d_model {self.d_model}")
+        x = x + self.pe[:, :seq_len]
+        return self.dropout(x)
+
+class LearnedPositionalEmbedding(nn.Module):
+    """Learned positional embeddings for transformer models."""
+    
+    def __init__(self, max_seq_len: int, d_model: int, dropout: float = 0.1):
+        super().__init__()
+        self.max_seq_len = max_seq_len
+        self.d_model = d_model
+        self.pos_embedding = nn.Embedding(max_seq_len, d_model)
+        self.dropout = nn.Dropout(dropout)
+        nn.init.normal_(self.pos_embedding.weight, std=0.02)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, d_model = x.size()
+        if seq_len > self.max_seq_len:
+            raise ValueError(f"Sequence length {seq_len} exceeds maximum {self.max_seq_len}")
+        if d_model != self.d_model:
+            raise ValueError(f"Input dimension {d_model} does not match d_model {self.d_model}")
+        positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(batch_size, -1)
+        pos_emb = self.pos_embedding(positions)
+        x = x + pos_emb
+        return self.dropout(x)
+
+class RotaryPositionalEmbedding(nn.Module):
+    """Rotary Positional Embedding (RoPE) for transformer models."""
+    
+    def __init__(self, d_model: int, max_seq_len: int = 2048, base: float = DEFAULT_BASE):
+        super().__init__()
+        self.d_model = d_model
+        self.max_seq_len = max_seq_len
+        self.base = base
+        inv_freq = 1.0 / (base ** (torch.arange(0, d_model, 2).float() / d_model))
+        self.register_buffer('inv_freq', inv_freq)
+        self._seq_len_cached = 0
+        self._cos_cached = None
+        self._sin_cached = None
+
+    def _update_cos_sin_cache(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> None:
+        if seq_len > self._seq_len_cached:
+            self._seq_len_cached = seq_len
+            t = torch.arange(seq_len, device=device, dtype=torch.float32)
+            freqs = torch.outer(t, self.inv_freq)
+            self._cos_cached = freqs.cos().to(dtype)
+            self._sin_cached = freqs.sin().to(dtype)
+
+    def _rotate_half(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        x1, x2 = x[..., :x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
+        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor, start_pos: int = 0) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size, seq_len, num_heads, head_dim = q.shape
+        self._update_cos_sin_cache(start_pos + seq_len, q.device, q.dtype)
+        cos = self._cos_cached[start_pos:start_pos + seq_len, :head_dim // 2].view(1, seq_len, 1, -1)
+        sin = self._sin_cached[start_pos:start_pos + seq_len, :head_dim // 2].view(1, seq_len, 1, -1)
+        q = q.transpose(1, 2).reshape(batch_size * num_heads, seq_len, head_dim)
+        k = k.transpose(1, 2).reshape(batch_size * num_heads, seq_len, head_dim)
+        q_rot = self._rotate_half(q, cos, sin)
+        k_rot = self._rotate_half(k, cos, sin)
+        q_rot = q_rot.reshape(batch_size, num_heads, seq_len, head_dim).transpose(1, 2)
+        k_rot = k_rot.reshape(batch_size, num_heads, seq_len, head_dim).transpose(1, 2)
+        return q_rot, k_rot
+
+class TechEmbeddingLayer(nn.Module):
+    """Comprehensive embedding layer with token and positional embeddings."""
+    
+    def __init__(
+        self,
+        vocab_size: int,
+        d_model: int,
+        max_seq_len: int = 512,
+        dropout: float = 0.1,
+        padding_idx: int = 0,
+        pos_encoding: str = "learned",
+        layer_norm: bool = True,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.vocab_size = vocab_size
+        self.padding_idx = padding_idx
+        self.pos_encoding_type = pos_encoding.lower()
+        self.token_embedding = nn.Embedding(vocab_size, d_model, padding_idx=padding_idx)
+        
+        if pos_encoding == "sinusoidal":
+            self.pos_encoding = PositionalEncoding(d_model, max_seq_len, dropout)
+        elif pos_encoding == "learned":
+            self.pos_encoding = LearnedPositionalEmbedding(max_seq_len, d_model, dropout)
+        elif pos_encoding == "rope":
+            self.pos_encoding = RotaryPositionalEmbedding(d_model, max_seq_len)
+        else:
+            raise ValueError(f"Unknown positional encoding type: {pos_encoding}")
+        
+        self.layer_norm = nn.LayerNorm(d_model) if layer_norm else nn.Identity()
+        self.dropout = nn.Dropout(dropout)
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02)
+        if self.padding_idx is not None:
+            nn.init.constant_(self.token_embedding.weight[self.padding_idx], 0.0)
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        if (input_ids >= self.vocab_size).any():
+            raise ValueError(f"Input IDs contain values >= vocab_size ({self.vocab_size})")
+        embeddings = self.token_embedding(input_ids)
+        if self.pos_encoding_type != "rope":
+            embeddings = self.pos_encoding(embeddings)
+        embeddings = self.layer_norm(embeddings)
+        return self.dropout(embeddings)
+
+    def get_positional_encoding(self) -> Optional[nn.Module]:
+        return self.pos_encoding if self.pos_encoding_type == "rope" else None
+
+def create_padding_mask(input_ids: torch.Tensor, padding_idx: int = 0) -> torch.Tensor:
+    return input_ids == padding_idx
+
+def create_causal_mask(seq_len: int, device: torch.device) -> torch.Tensor:
+    return torch.triu(torch.ones(seq_len, seq_len, device=device), diagonal=1).bool()
+
+# ============================================================================
+# MODEL CONSTANTS
 # ============================================================================
 
-# Constants for default configuration
 DEFAULT_D_MODEL = 512
 DEFAULT_N_HEADS = 8
 DEFAULT_N_LAYERS = 6
@@ -75,29 +214,20 @@ DEFAULT_PADDING_IDX = 0
 DEFAULT_ROUTER_TYPE = "adaptive"
 DEFAULT_VOCAB_SIZE = 10000
 
+# ============================================================================
+# MODEL COMPONENTS
+# ============================================================================
+
 class MultiHeadAttention(nn.Module):
     """Multi-head attention mechanism optimized for technical content."""
     
     def __init__(self, d_model: int, n_heads: int, dropout: float = DEFAULT_DROPOUT):
-        """
-        Initialize multi-head attention.
-
-        Args:
-            d_model (int): Dimension of the model embeddings.
-            n_heads (int): Number of attention heads.
-            dropout (float): Dropout rate for regularization.
-
-        Raises:
-            ValueError: If d_model is not divisible by n_heads.
-        """
         super().__init__()
         if d_model % n_heads != 0:
             raise ValueError(f"d_model ({d_model}) must be divisible by n_heads ({n_heads})")
-        
         self.d_model = d_model
         self.n_heads = n_heads
         self.d_k = d_model // n_heads
-        
         self.w_q = nn.Linear(d_model, d_model, bias=False)
         self.w_k = nn.Linear(d_model, d_model, bias=False)
         self.w_v = nn.Linear(d_model, d_model, bias=False)
@@ -106,7 +236,6 @@ class MultiHeadAttention(nn.Module):
         self._init_weights()
 
     def _init_weights(self) -> None:
-        """Initialize weights with Xavier uniform initialization."""
         for module in [self.w_q, self.w_k, self.w_v, self.w_o]:
             nn.init.xavier_uniform_(module.weight)
             if hasattr(module, 'bias') and module.bias is not None:
@@ -120,21 +249,7 @@ class MultiHeadAttention(nn.Module):
         mask: Optional[torch.Tensor] = None,
         pos_encoding: Optional[nn.Module] = None
     ) -> torch.Tensor:
-        """
-        Forward pass for multi-head attention.
-
-        Args:
-            query (torch.Tensor): Query tensor of shape (batch_size, seq_len, d_model).
-            key (torch.Tensor): Key tensor of shape (batch_size, seq_len, d_model).
-            value (torch.Tensor): Value tensor of shape (batch_size, seq_len, d_model).
-            mask (Optional[torch.Tensor]): Attention mask of shape (batch_size, seq_len, seq_len).
-            pos_encoding (Optional[nn.Module]): Positional encoding module (e.g., RoPE).
-
-        Returns:
-            torch.Tensor: Output tensor of shape (batch_size, seq_len, d_model).
-        """
         batch_size, seq_len, _ = query.size()
-        
         Q = self.w_q(query).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
         K = self.w_k(key).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
         V = self.w_v(value).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
@@ -158,34 +273,16 @@ class FeedForward(nn.Module):
     """Position-wise feed-forward network with GELU activation."""
     
     def __init__(self, d_model: int, dim_feedforward: int, dropout: float = DEFAULT_DROPOUT):
-        """
-        Initialize feed-forward network.
-
-        Args:
-            d_model (int): Dimension of the model embeddings.
-            dim_feedforward (int): Dimension of the feed-forward layer.
-            dropout (float): Dropout rate for regularization.
-        """
         super().__init__()
         self.linear1 = nn.Linear(d_model, dim_feedforward)
         self.linear2 = nn.Linear(dim_feedforward, d_model)
         self.dropout = nn.Dropout(dropout)
-        
         nn.init.xavier_uniform_(self.linear1.weight)
         nn.init.zeros_(self.linear1.bias)
         nn.init.xavier_uniform_(self.linear2.weight)
         nn.init.zeros_(self.linear2.bias)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Forward pass for feed-forward network.
-
-        Args:
-            x (torch.Tensor): Input tensor of shape (batch_size, seq_len, d_model).
-
-        Returns:
-            torch.Tensor: Output tensor of shape (batch_size, seq_len, d_model).
-        """
         x = F.gelu(self.linear1(x))
         x = self.dropout(x)
         return self.linear2(x)
@@ -194,17 +291,6 @@ class RecursionRouter(nn.Module):
     """Router to determine recursion steps for technical problem processing."""
     
     def __init__(self, d_model: int, max_steps: int = DEFAULT_MAX_STEPS, router_type: str = DEFAULT_ROUTER_TYPE):
-        """
-        Initialize recursion router.
-
-        Args:
-            d_model (int): Dimension of the model embeddings.
-            max_steps (int): Maximum number of recursion steps.
-            router_type (str): Type of router ('adaptive' or 'fixed').
-
-        Raises:
-            ValueError: If router_type is invalid.
-        """
         super().__init__()
         self.max_steps = max_steps
         self.router_type = router_type.lower()
@@ -223,15 +309,6 @@ class RecursionRouter(nn.Module):
             raise ValueError(f"Invalid router_type: {router_type}. Choose 'adaptive' or 'fixed'.")
 
     def forward(self, x: torch.Tensor) -> Union[torch.Tensor, int]:
-        """
-        Determine the number of recursion steps.
-
-        Args:
-            x (torch.Tensor): Input tensor of shape (batch_size, seq_len, d_model).
-
-        Returns:
-            Union[torch.Tensor, int]: Number of steps (tensor for adaptive, int for fixed).
-        """
         if self.router_type == "adaptive":
             seq_repr = x.mean(dim=1)
             step_probs = self.complexity_classifier(seq_repr)
@@ -250,21 +327,9 @@ class RecursiveTransformerLayer(nn.Module):
         dropout: float = DEFAULT_DROPOUT,
         router_type: str = DEFAULT_ROUTER_TYPE
     ):
-        """
-        Initialize recursive transformer layer.
-
-        Args:
-            d_model (int): Dimension of the model embeddings.
-            n_heads (int): Number of attention heads.
-            dim_feedforward (int): Dimension of the feed-forward layer.
-            max_steps (int): Maximum number of recursion steps.
-            dropout (float): Dropout rate for regularization.
-            router_type (str): Type of router ('adaptive' or 'fixed').
-        """
         super().__init__()
         self.max_steps = max_steps
         self.d_model = d_model
-        
         self.attention = MultiHeadAttention(d_model, n_heads, dropout)
         self.feedforward = FeedForward(d_model, dim_feedforward, dropout)
         self.norm1 = nn.LayerNorm(d_model)
@@ -274,7 +339,6 @@ class RecursiveTransformerLayer(nn.Module):
         self.step_projections = nn.ModuleList([
             nn.Linear(d_model, d_model) for _ in range(max_steps)
         ])
-        
         for proj in self.step_projections:
             nn.init.xavier_uniform_(proj.weight)
             nn.init.zeros_(proj.bias)
@@ -285,17 +349,6 @@ class RecursiveTransformerLayer(nn.Module):
         mask: Optional[torch.Tensor] = None,
         pos_encoding: Optional[nn.Module] = None
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Forward pass for recursive transformer layer.
-
-        Args:
-            x (torch.Tensor): Input tensor of shape (batch_size, seq_len, d_model).
-            mask (Optional[torch.Tensor]): Attention mask of shape (batch_size, seq_len, seq_len).
-            pos_encoding (Optional[nn.Module]): Positional encoding module (e.g., RoPE).
-
-        Returns:
-            Tuple[torch.Tensor, torch.Tensor]: Output tensor and computation loss.
-        """
         steps = self.router(x)
         if isinstance(steps, (int, torch.Tensor)) and not torch.is_tensor(steps):
             return self._recursive_forward_fixed(x, mask, steps, pos_encoding)
@@ -308,11 +361,9 @@ class RecursiveTransformerLayer(nn.Module):
         num_steps: int,
         pos_encoding: Optional[nn.Module]
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Fixed recursion forward pass."""
         device = x.device
         batch_size = x.shape[0]
         computation_loss = torch.tensor(0.0, device=device)
-        
         for step in range(min(num_steps, self.max_steps)):
             step_input = self.step_projections[step](x) if step < len(self.step_projections) else x
             attended = self.attention(step_input, step_input, step_input, mask, pos_encoding)
@@ -320,7 +371,6 @@ class RecursiveTransformerLayer(nn.Module):
             fed_forward = self.feedforward(x)
             x = self.norm2(x + self.dropout(fed_forward))
             computation_loss += torch.tensor(0.1, device=device) * batch_size
-        
         return x, computation_loss
 
     def _recursive_forward_adaptive(
@@ -330,18 +380,15 @@ class RecursiveTransformerLayer(nn.Module):
         steps: torch.Tensor,
         pos_encoding: Optional[nn.Module]
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Adaptive recursion forward pass."""
         batch_size, seq_len, d_model = x.shape
         device = x.device
         max_batch_steps = int(steps.max().item())
         computation_loss = torch.tensor(0.0, device=device)
         active_batches = torch.ones(batch_size, device=device, dtype=torch.bool)
-        
         for step in range(min(max_batch_steps, self.max_steps)):
             step_mask = (steps > step) & active_batches
             if not step_mask.any():
                 break
-            
             step_input = self.step_projections[step](x) if step < len(self.step_projections) else x
             attended = self.attention(step_input, step_input, step_input, mask, pos_encoding)
             attended = torch.where(step_mask.unsqueeze(-1).unsqueeze(-1), attended, torch.zeros_like(attended))
@@ -351,7 +398,6 @@ class RecursiveTransformerLayer(nn.Module):
             x = self.norm2(x + self.dropout(fed_forward))
             computation_loss += torch.tensor(0.1, device=device) * step_mask.sum()
             active_batches &= (steps > step)
-        
         return x, computation_loss
 
 class MixtureOfRecursions(nn.Module):
@@ -371,27 +417,10 @@ class MixtureOfRecursions(nn.Module):
         padding_idx: int = DEFAULT_PADDING_IDX,
         pos_encoding: str = "learned"
     ):
-        """
-        Initialize the Mixture of Recursions model.
-
-        Args:
-            vocab_size (int): Size of the vocabulary.
-            d_model (int): Dimension of the model embeddings.
-            n_layers (int): Number of transformer layers.
-            n_heads (int): Number of attention heads.
-            max_steps (int): Maximum number of recursion steps.
-            dim_feedforward (int): Dimension of the feed-forward layer.
-            dropout (float): Dropout rate for regularization.
-            max_seq_len (int): Maximum sequence length.
-            router_type (str): Type of router ('adaptive' or 'fixed').
-            padding_idx (int): Index for padding token.
-            pos_encoding (str): Type of positional encoding ('learned', 'sinusoidal', 'rope').
-        """
         super().__init__()
         self.d_model = d_model
         self.vocab_size = vocab_size
         self.padding_idx = padding_idx
-        
         self.embeddings = TechEmbeddingLayer(
             vocab_size=vocab_size,
             d_model=d_model,
@@ -400,7 +429,6 @@ class MixtureOfRecursions(nn.Module):
             padding_idx=padding_idx,
             pos_encoding=pos_encoding
         )
-        
         self.layers = nn.ModuleList([
             RecursiveTransformerLayer(
                 d_model=d_model,
@@ -411,39 +439,24 @@ class MixtureOfRecursions(nn.Module):
                 router_type=router_type
             ) for _ in range(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:
-        """Initialize weights for the language model head."""
         nn.init.xavier_uniform_(self.lm_head.weight)
 
     def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Forward pass for the model.
-
-        Args:
-            input_ids (torch.Tensor): Input tensor of shape (batch_size, seq_len).
-            attention_mask (Optional[torch.Tensor]): Attention mask of shape (batch_size, seq_len).
-
-        Returns:
-            Tuple[torch.Tensor, torch.Tensor]: Logits and total computation loss.
-        """
         batch_size, seq_len = input_ids.shape
         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)
-        
         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
@@ -455,27 +468,13 @@ class MixtureOfRecursions(nn.Module):
         top_k: Optional[int] = None,
         top_p: Optional[float] = None
     ) -> torch.Tensor:
-        """
-        Generate the next token for a given input sequence.
-
-        Args:
-            input_ids (torch.Tensor): Input tensor of shape (batch_size, seq_len).
-            temperature (float): Temperature for softmax scaling.
-            top_k (Optional[int]): Number of top-k tokens to sample from.
-            top_p (Optional[float]): Cumulative probability for nucleus sampling.
-
-        Returns:
-            torch.Tensor: Next token IDs of shape (batch_size, 1).
-        """
         self.eval()
         with torch.no_grad():
             logits, _ = self.forward(input_ids)
             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)
@@ -484,7 +483,6 @@ 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)
 
@@ -492,15 +490,6 @@ 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):
-        """
-        Initialize the text generator.
-
-        Args:
-            model (nn.Module): The transformer model.
-            tokenizer (Tokenizer): Tokenizer for encoding/decoding text.
-            max_length (int): Maximum sequence length for generation.
-            device (Optional[torch.device]): Device to run the model on.
-        """
         self.model = model
         self.tokenizer = tokenizer
         self.max_length = max_length
@@ -518,36 +507,16 @@ class TextGenerator:
         top_p: Optional[float] = 0.9,
         max_new_tokens: Optional[int] = None
     ) -> str:
-        """
-        Generate text based on a prompt.
-
-        Args:
-            prompt (str): Input prompt for generation.
-            method (str): Generation method ('greedy', 'sample', 'top_k', 'nucleus').
-            temperature (float): Temperature for softmax scaling.
-            top_k (Optional[int]): Number of top-k tokens to sample from.
-            top_p (Optional[float]): Cumulative probability for nucleus sampling.
-            max_new_tokens (Optional[int]): Maximum number of new tokens to generate.
-
-        Returns:
-            str: Generated text response.
-
-        Raises:
-            ValueError: If the generation method is invalid.
-        """
         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":
@@ -558,38 +527,30 @@ class TextGenerator:
                     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:
-        """Generate the next token using greedy decoding."""
         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:
-        """Generate the next token using random sampling."""
         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:
-        """Generate the next token using top-k sampling."""
         logits, _ = self.model(input_tensor)
         logits = logits[:, -1, :] / temperature
         top_k_logits, top_k_indices = torch.topk(logits, top_k)
@@ -598,19 +559,616 @@ class TextGenerator:
         return top_k_indices.gather(-1, next_token_idx)
 
     def _nucleus_generate(self, input_tensor: torch.Tensor, temperature: float, top_p: float) -> torch.Tensor:
-        """Generate the next token using nucleus (top-p) sampling."""
         return self.model.generate_step(input_tensor, temperature, top_p=top_p)
 
 def count_parameters(model: nn.Module) -> Tuple[int, int]:
-    """
-    Count total and trainable parameters in the model.
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    return total_params, trainable_params
+
+def main():
+    """Test the MixtureOfRecursions model and its components."""
+    print("Initializing MixtureOfRecursions model...")
+    model = MixtureOfRecursions(
+        vocab_size=DEFAULT_VOCAB_SIZE,
+        d_model=DEFAULT_D_MODEL,
+        n_layers=DEFAULT_N_LAYERS,
+        n_heads=DEFAULT_N_HEADS,
+        max_steps=DEFAULT_MAX_STEPS,
+        dim_feedforward=DEFAULT_DIM_FEEDFORWARD,
+        dropout=DEFAULT_DROPOUT,
+        router_type=DEFAULT_ROUTER_TYPE
+    )
+    
+    total_params, trainable_params = count_parameters(model)
+    print(f"Total parameters: {total_params:,}")
+    print(f"Trainable parameters: {trainable_params:,}")
+    
+    print("\nTesting forward pass...")
+    batch_size, seq_len = 4, 128
+    input_ids = torch.randint(0, DEFAULT_VOCAB_SIZE, (batch_size, seq_len))
+    attention_mask = torch.ones_like(input_ids)
+    attention_mask[:, -10:] = 0
+    
+    logits, comp_loss = model(input_ids, attention_mask)
+    
+    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)
+    print(f"Generated next token: {next_token.item()}")
+    
+    print("\nModel test completed successfully!")
 
-    Args:
-        model (nn.Module): The model to analyze.
+if __name__ == "__main__":
+    main()import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import Optional, Tuple, Union, List
+
+# ============================================================================
+# TRANSFORMERS COMPATIBILITY
+# ============================================================================
+from transformers import PretrainedConfig
+from transformers.modeling_utils import PreTrainedModel
+
+class MixtureOfRecursionsConfig(PretrainedConfig):
+    """Configuration class for MixtureOfRecursions model."""
+    
+    model_type = "mixture_of_recursions"
+    
+    def __init__(
+        self,
+        vocab_size=31985,
+        d_model=384,
+        n_layers=12,
+        n_heads=6,
+        max_steps=4,
+        dim_feedforward=2048,
+        dropout=0.1,
+        max_seq_len=128,
+        router_type="adaptive",
+        padding_idx=0,
+        pos_encoding="learned",
+        hidden_size=None,
+        num_hidden_layers=None,
+        num_attention_heads=None,
+        intermediate_size=None,
+        max_position_embeddings=None,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.max_steps = max_steps
+        self.dim_feedforward = dim_feedforward
+        self.dropout = dropout
+        self.max_seq_len = max_seq_len
+        self.router_type = router_type
+        self.padding_idx = padding_idx
+        self.pos_encoding = pos_encoding
+        self.hidden_size = hidden_size or d_model
+        self.num_hidden_layers = num_hidden_layers or n_layers
+        self.num_attention_heads = num_attention_heads or n_heads
+        self.intermediate_size = intermediate_size or dim_feedforward
+        self.max_position_embeddings = max_position_embeddings or max_seq_len
+
+# ============================================================================
+# EMBEDDINGS MODULE (merged from embeddings.py)
+# ============================================================================
+
+DEFAULT_BASE = 10000.0
+DEFAULT_CUTOFFS = [2000, 10000]
+DEFAULT_DIV_VAL = 4.0
+
+class PositionalEncoding(nn.Module):
+    """Sinusoidal positional encoding for transformer models."""
+    
+    def __init__(self, d_model: int, max_seq_len: int = 512, dropout: float = 0.1):
+        super().__init__()
+        self.d_model = d_model
+        self.dropout = nn.Dropout(dropout)
+        pe = torch.zeros(max_seq_len, d_model)
+        position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(DEFAULT_BASE) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term[:, :-1] if d_model % 2 == 1 else div_term)
+        self.register_buffer('pe', pe.unsqueeze(0))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, d_model = x.size()
+        if d_model != self.d_model:
+            raise ValueError(f"Input dimension {d_model} does not match d_model {self.d_model}")
+        x = x + self.pe[:, :seq_len]
+        return self.dropout(x)
+
+class LearnedPositionalEmbedding(nn.Module):
+    """Learned positional embeddings for transformer models."""
+    
+    def __init__(self, max_seq_len: int, d_model: int, dropout: float = 0.1):
+        super().__init__()
+        self.max_seq_len = max_seq_len
+        self.d_model = d_model
+        self.pos_embedding = nn.Embedding(max_seq_len, d_model)
+        self.dropout = nn.Dropout(dropout)
+        nn.init.normal_(self.pos_embedding.weight, std=0.02)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, d_model = x.size()
+        if seq_len > self.max_seq_len:
+            raise ValueError(f"Sequence length {seq_len} exceeds maximum {self.max_seq_len}")
+        if d_model != self.d_model:
+            raise ValueError(f"Input dimension {d_model} does not match d_model {self.d_model}")
+        positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(batch_size, -1)
+        pos_emb = self.pos_embedding(positions)
+        x = x + pos_emb
+        return self.dropout(x)
+
+class RotaryPositionalEmbedding(nn.Module):
+    """Rotary Positional Embedding (RoPE) for transformer models."""
+    
+    def __init__(self, d_model: int, max_seq_len: int = 2048, base: float = DEFAULT_BASE):
+        super().__init__()
+        self.d_model = d_model
+        self.max_seq_len = max_seq_len
+        self.base = base
+        inv_freq = 1.0 / (base ** (torch.arange(0, d_model, 2).float() / d_model))
+        self.register_buffer('inv_freq', inv_freq)
+        self._seq_len_cached = 0
+        self._cos_cached = None
+        self._sin_cached = None
+
+    def _update_cos_sin_cache(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> None:
+        if seq_len > self._seq_len_cached:
+            self._seq_len_cached = seq_len
+            t = torch.arange(seq_len, device=device, dtype=torch.float32)
+            freqs = torch.outer(t, self.inv_freq)
+            self._cos_cached = freqs.cos().to(dtype)
+            self._sin_cached = freqs.sin().to(dtype)
+
+    def _rotate_half(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        x1, x2 = x[..., :x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
+        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor, start_pos: int = 0) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size, seq_len, num_heads, head_dim = q.shape
+        self._update_cos_sin_cache(start_pos + seq_len, q.device, q.dtype)
+        cos = self._cos_cached[start_pos:start_pos + seq_len, :head_dim // 2].view(1, seq_len, 1, -1)
+        sin = self._sin_cached[start_pos:start_pos + seq_len, :head_dim // 2].view(1, seq_len, 1, -1)
+        q = q.transpose(1, 2).reshape(batch_size * num_heads, seq_len, head_dim)
+        k = k.transpose(1, 2).reshape(batch_size * num_heads, seq_len, head_dim)
+        q_rot = self._rotate_half(q, cos, sin)
+        k_rot = self._rotate_half(k, cos, sin)
+        q_rot = q_rot.reshape(batch_size, num_heads, seq_len, head_dim).transpose(1, 2)
+        k_rot = k_rot.reshape(batch_size, num_heads, seq_len, head_dim).transpose(1, 2)
+        return q_rot, k_rot
+
+class TechEmbeddingLayer(nn.Module):
+    """Comprehensive embedding layer with token and positional embeddings."""
+    
+    def __init__(
+        self,
+        vocab_size: int,
+        d_model: int,
+        max_seq_len: int = 512,
+        dropout: float = 0.1,
+        padding_idx: int = 0,
+        pos_encoding: str = "learned",
+        layer_norm: bool = True,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.vocab_size = vocab_size
+        self.padding_idx = padding_idx
+        self.pos_encoding_type = pos_encoding.lower()
+        self.token_embedding = nn.Embedding(vocab_size, d_model, padding_idx=padding_idx)
+        
+        if pos_encoding == "sinusoidal":
+            self.pos_encoding = PositionalEncoding(d_model, max_seq_len, dropout)
+        elif pos_encoding == "learned":
+            self.pos_encoding = LearnedPositionalEmbedding(max_seq_len, d_model, dropout)
+        elif pos_encoding == "rope":
+            self.pos_encoding = RotaryPositionalEmbedding(d_model, max_seq_len)
+        else:
+            raise ValueError(f"Unknown positional encoding type: {pos_encoding}")
+        
+        self.layer_norm = nn.LayerNorm(d_model) if layer_norm else nn.Identity()
+        self.dropout = nn.Dropout(dropout)
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02)
+        if self.padding_idx is not None:
+            nn.init.constant_(self.token_embedding.weight[self.padding_idx], 0.0)
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        if (input_ids >= self.vocab_size).any():
+            raise ValueError(f"Input IDs contain values >= vocab_size ({self.vocab_size})")
+        embeddings = self.token_embedding(input_ids)
+        if self.pos_encoding_type != "rope":
+            embeddings = self.pos_encoding(embeddings)
+        embeddings = self.layer_norm(embeddings)
+        return self.dropout(embeddings)
+
+    def get_positional_encoding(self) -> Optional[nn.Module]:
+        return self.pos_encoding if self.pos_encoding_type == "rope" else None
+
+def create_padding_mask(input_ids: torch.Tensor, padding_idx: int = 0) -> torch.Tensor:
+    return input_ids == padding_idx
+
+def create_causal_mask(seq_len: int, device: torch.device) -> torch.Tensor:
+    return torch.triu(torch.ones(seq_len, seq_len, device=device), diagonal=1).bool()
+
+# ============================================================================
+# MODEL CONSTANTS
+# ============================================================================
+
+DEFAULT_D_MODEL = 512
+DEFAULT_N_HEADS = 8
+DEFAULT_N_LAYERS = 6
+DEFAULT_MAX_STEPS = 4
+DEFAULT_DIM_FEEDFORWARD = 2048
+DEFAULT_DROPOUT = 0.1
+DEFAULT_MAX_SEQ_LEN = 512
+DEFAULT_PADDING_IDX = 0
+DEFAULT_ROUTER_TYPE = "adaptive"
+DEFAULT_VOCAB_SIZE = 10000
+
+# ============================================================================
+# MODEL COMPONENTS
+# ============================================================================
 
-    Returns:
-        Tuple[int, int]: Total and trainable parameter counts.
-    """
+class MultiHeadAttention(nn.Module):
+    """Multi-head attention mechanism optimized for technical content."""
+    
+    def __init__(self, d_model: int, n_heads: int, dropout: float = DEFAULT_DROPOUT):
+        super().__init__()
+        if d_model % n_heads != 0:
+            raise ValueError(f"d_model ({d_model}) must be divisible by n_heads ({n_heads})")
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.d_k = d_model // n_heads
+        self.w_q = nn.Linear(d_model, d_model, bias=False)
+        self.w_k = nn.Linear(d_model, d_model, bias=False)
+        self.w_v = nn.Linear(d_model, d_model, bias=False)
+        self.w_o = nn.Linear(d_model, d_model)
+        self.dropout = nn.Dropout(dropout)
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        for module in [self.w_q, self.w_k, self.w_v, self.w_o]:
+            nn.init.xavier_uniform_(module.weight)
+            if hasattr(module, 'bias') and module.bias is not None:
+                nn.init.zeros_(module.bias)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        pos_encoding: Optional[nn.Module] = None
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = query.size()
+        Q = self.w_q(query).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
+        K = self.w_k(key).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
+        V = self.w_v(value).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
+        
+        if pos_encoding is not None:
+            Q, K = pos_encoding(Q, K)
+        
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
+        
+        if mask is not None:
+            mask = mask.unsqueeze(1).expand(batch_size, self.n_heads, seq_len, seq_len)
+            scores = scores.masked_fill(mask, float('-inf'))
+        
+        attention_weights = F.softmax(scores, dim=-1)
+        attention_weights = self.dropout(attention_weights)
+        attended = torch.matmul(attention_weights, V)
+        attended = attended.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
+        return self.w_o(attended)
+
+class FeedForward(nn.Module):
+    """Position-wise feed-forward network with GELU activation."""
+    
+    def __init__(self, d_model: int, dim_feedforward: int, dropout: float = DEFAULT_DROPOUT):
+        super().__init__()
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+        self.dropout = nn.Dropout(dropout)
+        nn.init.xavier_uniform_(self.linear1.weight)
+        nn.init.zeros_(self.linear1.bias)
+        nn.init.xavier_uniform_(self.linear2.weight)
+        nn.init.zeros_(self.linear2.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.gelu(self.linear1(x))
+        x = self.dropout(x)
+        return self.linear2(x)
+
+class RecursionRouter(nn.Module):
+    """Router to determine recursion steps for technical problem processing."""
+    
+    def __init__(self, d_model: int, max_steps: int = DEFAULT_MAX_STEPS, router_type: str = DEFAULT_ROUTER_TYPE):
+        super().__init__()
+        self.max_steps = max_steps
+        self.router_type = router_type.lower()
+        
+        if self.router_type == "adaptive":
+            self.complexity_classifier = nn.Sequential(
+                nn.Linear(d_model, d_model // 4),
+                nn.GELU(),
+                nn.Dropout(DEFAULT_DROPOUT),
+                nn.Linear(d_model // 4, max_steps + 1),
+                nn.Softmax(dim=-1)
+            )
+        elif self.router_type == "fixed":
+            self.register_buffer('fixed_steps', torch.tensor(max_steps, dtype=torch.long))
+        else:
+            raise ValueError(f"Invalid router_type: {router_type}. Choose 'adaptive' or 'fixed'.")
+
+    def forward(self, x: torch.Tensor) -> Union[torch.Tensor, int]:
+        if self.router_type == "adaptive":
+            seq_repr = x.mean(dim=1)
+            step_probs = self.complexity_classifier(seq_repr)
+            return torch.argmax(step_probs, dim=-1)
+        return self.fixed_steps.item()
+
+class RecursiveTransformerLayer(nn.Module):
+    """Transformer layer with recursive computation capability."""
+    
+    def __init__(
+        self,
+        d_model: int,
+        n_heads: int,
+        dim_feedforward: int,
+        max_steps: int = DEFAULT_MAX_STEPS,
+        dropout: float = DEFAULT_DROPOUT,
+        router_type: str = DEFAULT_ROUTER_TYPE
+    ):
+        super().__init__()
+        self.max_steps = max_steps
+        self.d_model = d_model
+        self.attention = MultiHeadAttention(d_model, n_heads, dropout)
+        self.feedforward = FeedForward(d_model, dim_feedforward, dropout)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.router = RecursionRouter(d_model, max_steps, router_type)
+        self.step_projections = nn.ModuleList([
+            nn.Linear(d_model, d_model) for _ in range(max_steps)
+        ])
+        for proj in self.step_projections:
+            nn.init.xavier_uniform_(proj.weight)
+            nn.init.zeros_(proj.bias)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        pos_encoding: Optional[nn.Module] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        steps = self.router(x)
+        if isinstance(steps, (int, torch.Tensor)) and not torch.is_tensor(steps):
+            return self._recursive_forward_fixed(x, mask, steps, pos_encoding)
+        return self._recursive_forward_adaptive(x, mask, steps, pos_encoding)
+
+    def _recursive_forward_fixed(
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        num_steps: int,
+        pos_encoding: Optional[nn.Module]
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        device = x.device
+        batch_size = x.shape[0]
+        computation_loss = torch.tensor(0.0, device=device)
+        for step in range(min(num_steps, self.max_steps)):
+            step_input = self.step_projections[step](x) if step < len(self.step_projections) else x
+            attended = self.attention(step_input, step_input, step_input, mask, pos_encoding)
+            x = self.norm1(x + self.dropout(attended))
+            fed_forward = self.feedforward(x)
+            x = self.norm2(x + self.dropout(fed_forward))
+            computation_loss += torch.tensor(0.1, device=device) * batch_size
+        return x, computation_loss
+
+    def _recursive_forward_adaptive(
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        steps: torch.Tensor,
+        pos_encoding: Optional[nn.Module]
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size, seq_len, d_model = x.shape
+        device = x.device
+        max_batch_steps = int(steps.max().item())
+        computation_loss = torch.tensor(0.0, device=device)
+        active_batches = torch.ones(batch_size, device=device, dtype=torch.bool)
+        for step in range(min(max_batch_steps, self.max_steps)):
+            step_mask = (steps > step) & active_batches
+            if not step_mask.any():
+                break
+            step_input = self.step_projections[step](x) if step < len(self.step_projections) else x
+            attended = self.attention(step_input, step_input, step_input, mask, pos_encoding)
+            attended = torch.where(step_mask.unsqueeze(-1).unsqueeze(-1), attended, torch.zeros_like(attended))
+            x = self.norm1(x + self.dropout(attended))
+            fed_forward = self.feedforward(x)
+            fed_forward = torch.where(step_mask.unsqueeze(-1).unsqueeze(-1), fed_forward, torch.zeros_like(fed_forward))
+            x = self.norm2(x + self.dropout(fed_forward))
+            computation_loss += torch.tensor(0.1, device=device) * step_mask.sum()
+            active_batches &= (steps > step)
+        return x, computation_loss
+
+class MixtureOfRecursions(nn.Module):
+    """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
+        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
+        )
+        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)
+        ])
+        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)
+
+    def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size, seq_len = input_ids.shape
+        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)
+        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
+
+    def generate_step(
+        self,
+        input_ids: torch.Tensor,
+        temperature: float = 1.0,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None
+    ) -> torch.Tensor:
+        self.eval()
+        with torch.no_grad():
+            logits, _ = self.forward(input_ids)
+            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)
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+                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)
+
+def count_parameters(model: nn.Module) -> Tuple[int, int]:
     total_params = sum(p.numel() for p in model.parameters())
     trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
     return total_params, trainable_params