LLaMA Components › Stacking Blocks into a Model

class LLaMAModel:
    """
    Simplified LLaMA model for demonstration.

    Stacks transformer blocks with token embeddings and output projection.
    """

    def __init__(
        self, vocab_size, d_model, n_heads, d_ff, n_layers, rope_base=10000
    ):
        self.d_model = d_model

        # Token embeddings (no positional embeddings - RoPE handles position)
        self.token_embed = np.random.randn(vocab_size, d_model) * 0.02

        # Stack of transformer blocks
        self.blocks = [
            LLaMABlock(d_model, n_heads, d_ff, rope_base)
            for _ in range(n_layers)
        ]

        # Final normalization
        self.final_norm = RMSNorm(d_model)

        # Output projection (weight-tied with embedding in real LLaMA)
        self.output_proj = np.random.randn(d_model, vocab_size) * 0.02

    def __call__(self, token_ids):
        """
        Forward pass.

        Args:
            token_ids: Array of token IDs, shape (seq_len,)

        Returns:
            Logits of shape (seq_len, vocab_size)
        """
        seq_len = len(token_ids)

        # Token embeddings
        x = self.token_embed[token_ids]  # (seq_len, d_model)

        # Causal mask
        mask = np.tril(np.ones((seq_len, seq_len), dtype=bool))

        # Pass through all blocks
        for block in self.blocks:
            x = block(x, mask)

        # Final normalization and projection
        x = self.final_norm(x)
        logits = x @ self.output_proj

        return logits

# Create a small model
model = LLaMAModel(
    vocab_size=1000,
    d_model=128,
    n_heads=4,
    d_ff=341,  # ~2.67 * 128
    n_layers=4,
)

# Sample input (random token IDs)
tokens = np.array([42, 123, 7, 891, 456, 0, 333, 222])
logits = model(tokens)