Most engineers are still writing prompts.
Power up your AI engineering skills.

# ==========================================================================
# Step 13: Create simple n-gram baseline for comparison
# ==========================================================================
# Baseline Model for Comparison
# We create a simpler model without attention to demonstrate the value
# that transformer blocks add to the basic n-gram foundation.
 
class SimpleNgramModel(nn.Module):
    """
    Simple n-gram model without attention.
 
    This model only uses character embeddings and position embeddings
    without any attention mechanisms. It serves as a baseline to show
    the improvement that transformers provide.
    """
 
    def __init__(self, vocab_size, n_embd=128, block_size=64):
        """
        Initialize the simple n-gram model.
 
        Args:
            vocab_size: Number of unique characters
            n_embd: Embedding dimension
            block_size: Maximum sequence length
        """
        super().__init__()
        self.vocab_size = vocab_size
        self.block_size = block_size
        # Only basic embeddings - no attention mechanisms
        self.token_embedding = nn.Embedding(vocab_size, n_embd)    # Character embeddings
        self.position_embedding = nn.Embedding(block_size, n_embd) # Position embeddings
        self.ln = nn.LayerNorm(n_embd)                             # Layer normalization
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)   # Output projection
 
    def forward(self, idx, targets=None):
        """
        Forward pass through the simple n-gram model.
 
        This is much simpler than the hybrid model - just embeddings
        and a direct projection to vocabulary, with no attention.
        """
        B, T = idx.size()
        pos = torch.arange(0, T, dtype=torch.long, device=idx.device)
 
        # Simple embedding combination (no attention)
        tok_emb = self.token_embedding(idx)     # Character embeddings
        pos_emb = self.position_embedding(pos)  # Position embeddings
        x = self.ln(tok_emb + pos_emb)          # Combine and normalize
        logits = self.lm_head(x)                # Project to vocabulary
 
        # Calculate loss if targets provided
        loss = None
        if targets is not None:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)
 
        return logits, loss
 
# Create the baseline model
ngram_model = SimpleNgramModel(vocab_size, config['n_embd'], config['block_size'])