Update model_slm.py

model_slm.py

@@ -2,86 +2,211 @@ 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
+
+# Constants for default configuration
+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
+
 class MultiHeadAttention(nn.Module):
-    """Multi-head attention mechanism optimized for technical content"""    
-    def __init__(self, d_model, n_heads, dropout=0.1):
-        super(MultiHeadAttention, self).__init__()
-        assert d_model % n_heads == 0        
+    """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.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.w_o = nn.Linear(d_model, d_model)
         self.dropout = nn.Dropout(dropout)
-        self._init_weights()    
-    def _init_weights(self):
-        """Initialize weights with Xavier uniform"""
+        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)    
-    def forward(self, query, key, value, mask=None, pos_encoding=None):
-        batch_size, seq_len, d_model = query.size()        
+            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:
+        """
+        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)        
+        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)        
+            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.masked_fill_(mask, float('-inf'))        
+            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, d_model)
-        output = self.w_o(attended)        
-        return output
+        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, dim_feedforward, dropout=0.1):
-        super(FeedForward, self).__init__()
+    """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)        
+        self.dropout = nn.Dropout(dropout)
+        
         nn.init.xavier_uniform_(self.linear1.weight)
-        nn.init.xavier_uniform_(self.linear2.weight)    
-    def forward(self, x):
+        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)
-        x = self.linear2(x)
-        return x
+        return self.linear2(x)
+
 class RecursionRouter(nn.Module):
-    """Router to decide recursion steps for different types of technical problems"""    
-    def __init__(self, d_model, max_steps=4, router_type="adaptive"):
-        super(RecursionRouter, self).__init__()
+    """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        
-        if router_type == "adaptive":
+        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(0.1),
+                nn.Dropout(DEFAULT_DROPOUT),
                 nn.Linear(d_model // 4, max_steps + 1),
                 nn.Softmax(dim=-1)
             )
-        elif router_type == "fixed":
-            self.fixed_steps = max_steps    
-    def forward(self, x):
+        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]:
+        """
+        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)
-            steps = torch.argmax(step_probs, dim=-1)
-            return steps
-        return self.fixed_steps
+            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, n_heads, dim_feedforward, max_steps=4, 
-                 dropout=0.1, router_type="adaptive"):
-        super(RecursiveTransformerLayer, self).__init__()        
+    """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
+    ):
+        """
+        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.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)
@@ -90,35 +215,75 @@ class RecursiveTransformerLayer(nn.Module):
         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)
-        ])    
-    def forward(self, x, mask=None, pos_encoding=None):
+        ])
+        
+        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]:
+        """
+        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):
-            num_steps = min(steps, self.max_steps)
-            return self._recursive_forward_fixed(x, mask, num_steps, pos_encoding)
-        return self._recursive_forward_adaptive(x, mask, steps, pos_encoding)    
-    def _recursive_forward_fixed(self, x, mask, num_steps, pos_encoding):
+        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]:
+        """Fixed recursion forward pass."""
         device = x.device
         batch_size = x.shape[0]
-        computation_loss = torch.tensor(0.0, device=device)        
-        for step in range(num_steps):
+        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, mask, steps, pos_encoding):
+            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]:
+        """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)        
+        computation_loss = torch.tensor(0.0, device=device)
         active_batches = torch.ones(batch_size, device=device, dtype=torch.bool)
-        for step in range(max_batch_steps):
+        
+        for step in range(min(max_batch_steps, self.max_steps)):
             step_mask = (steps > step) & active_batches
             if not step_mask.any():
-                break                
+                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))
@@ -127,25 +292,57 @@ class RecursiveTransformerLayer(nn.Module):
             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)        
+            active_batches &= (steps > step)
+        
         return x, computation_loss
+
 class MixtureOfRecursions(nn.Module):
-    """Main model with mixture of recursive transformer layers"""    
-    def __init__(self, vocab_size, d_model=512, n_layers=6, n_heads=8, 
-                 max_steps=4, dim_feedforward=2048, dropout=0.1, 
-                 max_seq_len=512, router_type="adaptive", padding_idx=0):
-        super(MixtureOfRecursions, self).__init__()        
+    """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"
+    ):
+        """
+        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.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="learned"
-        )        
+            pos_encoding=pos_encoding
+        )
+        
         self.layers = nn.ModuleList([
             RecursiveTransformerLayer(
                 d_model=d_model,
@@ -155,70 +352,144 @@ class MixtureOfRecursions(nn.Module):
                 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):
-        nn.init.xavier_uniform_(self.lm_head.weight)    
-    def forward(self, input_ids, attention_mask=None):
-        batch_size, seq_len = input_ids.shape        
+        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)
-        padding_mask = padding_mask.unsqueeze(1).expand(batch_size, seq_len, seq_len)
-        combined_mask = padding_mask | causal_mask.unsqueeze(0)        
+        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()        
-        device = x.device
-        total_computation_loss = torch.tensor(0.0, device=device)
+        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        
+            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, temperature=1.0, top_k=None, top_p=None):
+        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:
+        """
+        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            
+            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[indices_to_remove] = float('-inf')            
+                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] = 0
+                sorted_indices_to_remove[..., 0] = False
                 indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
-                last_logits[indices_to_remove] = float('-inf')            
+                last_logits = last_logits.masked_fill(indices_to_remove, float('-inf'))
+            
             probs = F.softmax(last_logits, dim=-1)
-            next_token = torch.multinomial(probs, num_samples=1)
-            return next_token
+            return torch.multinomial(probs, num_samples=1)
+
 class TextGenerator:
-    """Text generation utility for the tech model"""    
-    def __init__(self, model, tokenizer, max_length=100, device=None):
+    """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
         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, method="nucleus", temperature=1.0, top_k=50, top_p=0.9, max_new_tokens=None):
-        if max_new_tokens is None:
-            max_new_tokens = self.max_length        
+        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:
+        """
+        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)        
+        input_tensor = torch.tensor([input_ids], device=self.device)
+        
         self.model.eval()
-        generated_ids = []        
+        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:]                
-                # Generate next token
+                    input_tensor = input_tensor[:, -self.max_length:]
+                
                 if method == "greedy":
                     next_token = self._greedy_generate(input_tensor)
                 elif method == "sample":
@@ -228,75 +499,101 @@ class TextGenerator:
                 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}")                
+                    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        
-        # Decode the full sequence
+                    break
+        
         full_ids = input_ids + generated_ids
-        full_text = self.tokenizer.decode_ids(full_ids, skip_special_tokens=False)        
-        # Extract assistant response
+        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):
+            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, temperature):
+        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, temperature, top_k):
+        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)
         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, temperature, top_p):
+        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):
+
+def count_parameters(model: nn.Module) -> Tuple[int, int]:
+    """
+    Count total and trainable parameters in the model.
+
+    Args:
+        model (nn.Module): The model to analyze.
+
+    Returns:
+        Tuple[int, int]: Total and trainable parameter counts.
+    """
     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():
-    vocab_size = 10000
-    d_model = 512
-    n_layers = 6
-    n_heads = 8
-    seq_len = 128
-    batch_size = 4    
+    """Test the MixtureOfRecursions model and its components."""
     print("Initializing MixtureOfRecursions model...")
     model = MixtureOfRecursions(
-        vocab_size=vocab_size,
-        d_model=d_model,
-        n_layers=n_layers,
-        n_heads=n_heads,
-        max_steps=4,
-        dim_feedforward=2048,
-        dropout=0.1,
-        router_type="adaptive"
-    )    
+        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(f"Trainable parameters: {trainable_params:,}")
+    
     print("\nTesting forward pass...")
-    input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
+    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    
+    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}")
-    logits, comp_loss = model(input_ids, attention_mask)    
     print(f"Output logits shape: {logits.shape}")
-    print(f"Computation loss: {comp_loss}")
-    print(f"Expected logits shape: ({batch_size}, {seq_len}, {vocab_size})")    
+    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}")    
+    print(f"Generated next token: {next_token.item()}")
+    
     print("\nModel test completed successfully!")
+
 if __name__ == "__main__":
     main()