Class BertModel

models/bert_base.py:69–114

BERT-base: hidden_size=768, num_layers=12, num_heads=12, intermediate=3072 (110M params).

class BertModel(nn.Module):
    """
    BERT-base: hidden_size=768, num_layers=12, num_heads=12, intermediate=3072 (110M params).
    """

    def __init__(
        self,
        vocab_size: int = 30522,
        hidden_size: int = 768,
        num_layers: int = 12,
        num_heads: int = 12,
        intermediate_size: int = 3072,
        max_seq_len: int = 512,
        dropout: float = 0.0,
    ):
        super().__init__()
        self.word_embeddings = nn.Embedding(vocab_size, hidden_size)
        self.position_embeddings = nn.Embedding(max_seq_len, hidden_size)
        self.token_type_embeddings = nn.Embedding(2, hidden_size)
        self.embed_norm = nn.LayerNorm(hidden_size)
        self.embed_dropout = nn.Dropout(dropout)

        self.layers = nn.ModuleList([
            BertLayer(hidden_size, num_heads, intermediate_size, dropout)
            for _ in range(num_layers)
        ])

        self.pooler = nn.Linear(hidden_size, hidden_size)

        n_params = sum(p.numel() for p in self.parameters())
        print(f"BertModel: {n_params / 1e6:.1f}M parameters")

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        B, T = input_ids.shape
        positions = torch.arange(T, device=input_ids.device).unsqueeze(0)
        token_types = torch.zeros_like(input_ids)

        x = self.word_embeddings(input_ids) + self.position_embeddings(positions) + self.token_type_embeddings(token_types)
        x = self.embed_dropout(self.embed_norm(x))

        for layer in self.layers:
            x = layer(x)

        # Pooled output from [CLS] token
        pooled = torch.tanh(self.pooler(x[:, 0]))
        return pooled