| from __future__ import annotations | |
| from typing import Optional | |
| from transformers import PretrainedConfig | |
| from babylm.config.schema import ModelConfig | |
| class ModernBertSmallConfig(PretrainedConfig): | |
| """HF-compatible config for the from-scratch ModernBERT-Small architecture. | |
| Bridges to the pydantic `ModelConfig` (used for JSON validation/CLI) via `to_hf_config`, | |
| while giving `save_pretrained`/`from_pretrained`/`AutoConfig` compatibility. | |
| """ | |
| model_type = "modernbert_small" | |
| def __init__( | |
| self, | |
| vocab_size: int = 30522, | |
| hidden_size: int = 384, | |
| num_hidden_layers: int = 16, | |
| num_attention_heads: int = 6, | |
| intermediate_size: int = 576, | |
| hidden_activation: str = "gelu", | |
| attention_pattern: str = "alternating", | |
| global_attn_every_n_layers: int = 3, | |
| global_rope_theta: float = 160000.0, | |
| local_rope_theta: float = 10000.0, | |
| local_attention_window: int = 128, | |
| norm_eps: float = 1e-5, | |
| attention_dropout: float = 0.0, | |
| mlp_dropout: float = 0.0, | |
| embedding_dropout: float = 0.0, | |
| embedding_type: str = "standard", | |
| factorized_linear_bottleneck_dim: int = 128, | |
| factorized_mlp_bottleneck_dim: int = 128, | |
| initializer_range: float = 0.02, | |
| pad_token_id: int = 0, | |
| tie_word_embeddings: bool = True, | |
| **kwargs, | |
| ) -> None: | |
| super().__init__(pad_token_id=pad_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs) | |
| self.vocab_size = vocab_size | |
| self.hidden_size = hidden_size | |
| self.num_hidden_layers = num_hidden_layers | |
| self.num_attention_heads = num_attention_heads | |
| self.intermediate_size = intermediate_size | |
| self.hidden_activation = hidden_activation | |
| self.attention_pattern = attention_pattern | |
| self.global_attn_every_n_layers = global_attn_every_n_layers | |
| self.global_rope_theta = global_rope_theta | |
| self.local_rope_theta = local_rope_theta | |
| self.local_attention_window = local_attention_window | |
| self.norm_eps = norm_eps | |
| self.attention_dropout = attention_dropout | |
| self.mlp_dropout = mlp_dropout | |
| self.embedding_dropout = embedding_dropout | |
| self.embedding_type = embedding_type | |
| self.factorized_linear_bottleneck_dim = factorized_linear_bottleneck_dim | |
| self.factorized_mlp_bottleneck_dim = factorized_mlp_bottleneck_dim | |
| self.initializer_range = initializer_range | |
| def to_hf_config(model_cfg: ModelConfig, pad_token_id: Optional[int] = None) -> ModernBertSmallConfig: | |
| """Builds a ModernBertSmallConfig from the pydantic ModelConfig, optionally overriding | |
| pad_token_id with the value discovered from a trained tokenizer.""" | |
| data = model_cfg.model_dump() | |
| if pad_token_id is not None: | |
| data["pad_token_id"] = pad_token_id | |
| return ModernBertSmallConfig(**data) | |