Encoder Architecture › Complete Encoder Implementation

class BertEncoder:
    """
    Complete BERT-style encoder implementation.

    Includes embeddings, stacked encoder layers, and pooling.
    """

    def __init__(
        self,
        vocab_size=30522,
        max_position=512,
        n_layers=12,
        d_model=768,
        n_heads=12,
        d_ff=3072,
    ):
        self.d_model = d_model
        self.n_layers = n_layers

        # Embeddings
        self.token_embeddings = np.random.randn(vocab_size, d_model) * 0.02
        self.position_embeddings = np.random.randn(max_position, d_model) * 0.02
        self.segment_embeddings = np.random.randn(2, d_model) * 0.02
        self.embedding_norm = LayerNorm(d_model)

        # Encoder layers
        self.layers = [
            EncoderLayer(d_model, n_heads, d_ff) for _ in range(n_layers)
        ]

        # Final normalization
        self.final_norm = LayerNorm(d_model)

        # Pooler (for [CLS] representation)
        self.pooler_dense = np.random.randn(d_model, d_model) * 0.02

    def embed(self, token_ids, segment_ids=None):
        """
        Compute input embeddings.

        Args:
            token_ids: Array of token indices, shape (seq_len,)
            segment_ids: Optional segment indices for sentence pairs

        Returns:
            Embeddings of shape (seq_len, d_model)
        """
        seq_len = len(token_ids)

        # Token embeddings
        token_embeds = self.token_embeddings[token_ids]

        # Position embeddings
        positions = np.arange(seq_len)
        position_embeds = self.position_embeddings[positions]

        # Segment embeddings (default to segment 0)
        if segment_ids is None:
            segment_ids = np.zeros(seq_len, dtype=int)
        segment_embeds = self.segment_embeddings[segment_ids]

        # Sum and normalize
        embeddings = token_embeds + position_embeds + segment_embeds
        embeddings = self.embedding_norm(embeddings)

        return embeddings

    def __call__(self, token_ids, segment_ids=None, attention_mask=None):
        """
        Full forward pass.

        Args:
            token_ids: Token indices, shape (seq_len,)
            segment_ids: Segment indices for sentence pairs
            attention_mask: Padding mask (1 for real, 0 for padding)

        Returns:
            Dictionary with sequence_output and pooled_output
        """
        # Compute embeddings
        hidden_states = self.embed(token_ids, segment_ids)

        # Convert attention mask to additive format
        if attention_mask is not None:
            # attention_mask: 1 for real tokens, 0 for padding
            # Convert to: 0 for allowed, -inf for blocked
            extended_mask = (1.0 - attention_mask) * float("-inf")
            extended_mask = extended_mask[
                np.newaxis, np.newaxis, :
            ]  # Broadcast dims
        else:
            extended_mask = None

        # Pass through encoder layers
        for layer in self.layers:
            hidden_states, _ = layer(hidden_states)

        # Final normalization
        sequence_output = self.final_norm(hidden_states)

        # Pooled output (from [CLS] token with tanh activation)
        cls_output = sequence_output[0]
        pooled_output = np.tanh(cls_output @ self.pooler_dense)

        return {
            "last_hidden_state": sequence_output,
            "pooler_output": pooled_output,
        }


# Create a small encoder for demonstration
small_encoder = BertEncoder(
    vocab_size=1000,
    max_position=128,
    n_layers=4,
    d_model=256,
    n_heads=8,
    d_ff=1024,
)

# Process a sample sequence
sample_tokens = np.array([2, 105, 234, 567, 89, 3])  # [CLS], tokens..., [SEP]
output = small_encoder(sample_tokens)