"""HuggingFace PretrainedConfig for the dual-attention (DAT) decoder LM. This is the source-of-truth copy. scripts/convert_dat_to_hf.py copies it into a generated HF repository (alongside modeling_dat.py and the flattened model source) so the model can be loaded with ``AutoModelForCausalLM.from_pretrained(path, trust_remote_code=True)``. The config simply carries every field of models.attention.dat.config.DatLMConfig so that modeling_dat.py can rebuild the exact DatLMConfig at load time. """ from transformers import PretrainedConfig # Every field needed to rebuild DatLMConfig. New exported fields are appended # after existing constructor parameters to preserve positional compatibility. # modeling_dat.py rebuilds DatLMConfig as # DatLMConfig(**{f: getattr(config, f) for f in DAT_LM_FIELDS}). DAT_LM_FIELDS = ( "vocab_size", "max_seq_len", "pe_type", "hidden_dim", "n_heads_sa", "n_heads_ra", "n_layers", "dropout", "dff_factor", "ffn_hidden_dim_mode", "ffn_activation", "rope_theta", "max_rel_pos", "init_range", "init_scheme", "norm_type", "norm_first", "use_bias_qkv", "use_bias_out", "use_bias_ffn", "tie_lm_head", "symbol_dim", "n_symbols", "symbolic_attn_n_heads", "symbol_retrieval", "symbolic_use_bias", "shared_symbol_retriever", "share_attn_params", "positional_symbols_sinusoidal", "relative_symbols_rope", "relsymbolic_rel_n_heads", "relsymbolic_symbolic_attn_n_heads", "relsymbolic_neighborhood_size", "relsymbolic_include_self", "relsymbolic_normalize_rels", "relsymbolic_trainable_symbols", "relsymbolic_dropout", "relsymbolic_rel_scale", "relsymbolic_symbolic_attn_scale", "relsymbolic_use_bias", "ra_type", "ra_n_relations", "ra_rel_activation", "ra_symmetric_rels", "sequence_boundary_policy", "segment_boundary_token_id", "pad_token_id", "bos_token_id", "eos_token_id", "mlm_head_enabled", ) class DatConfig(PretrainedConfig): model_type = "dat" # Expose the model width under HF's conventional name so generic eval # harnesses (e.g. babylm-eval's finetuning head, which reads # config.hidden_size) can size their classification head. The model's own # field is hidden_dim; this maps the alias onto it for both get and set. attribute_map = {"hidden_size": "hidden_dim"} def __init__( self, vocab_size: int = 16384, max_seq_len: int = 513, pe_type: str = "rope", hidden_dim: int = 256, n_heads_sa: int = 2, n_heads_ra: int = 2, n_layers: int = 4, dropout: float = 0.0, dff_factor: int = 4, ffn_hidden_dim_mode: str = "dff_factor", ffn_activation: str = "gelu", rope_theta: float = 10000.0, max_rel_pos: int | None = None, init_range: float = 0.15, init_scheme: str = "xavier_uniform", norm_type: str = "rmsnorm", norm_first: bool = True, use_bias_qkv: bool = False, use_bias_out: bool = True, use_bias_ffn: bool = True, tie_lm_head: bool = True, symbol_dim: int | None = None, n_symbols: int | None = None, symbolic_attn_n_heads: int | None = None, symbol_retrieval: str = "symbolic", symbolic_use_bias: bool = False, shared_symbol_retriever: bool = True, share_attn_params: bool = False, positional_symbols_sinusoidal: bool = False, relative_symbols_rope: bool = False, relsymbolic_rel_n_heads: int = 4, relsymbolic_symbolic_attn_n_heads: int = 4, relsymbolic_neighborhood_size: int = 2, relsymbolic_include_self: bool = False, relsymbolic_normalize_rels: bool = True, relsymbolic_trainable_symbols: bool = True, relsymbolic_dropout: float = 0.0, relsymbolic_rel_scale: float | None = None, relsymbolic_symbolic_attn_scale: float | None = None, relsymbolic_use_bias: bool = False, ra_type: str = "ra", ra_n_relations: int | None = None, ra_rel_activation: str = "identity", ra_symmetric_rels: bool = False, sequence_boundary_policy: str = "eos_document", segment_boundary_token_id: int | None = None, pad_token_id: int = 0, bos_token_id: int = 1, eos_token_id: int = 2, mlm_head_enabled: bool = False, **kwargs, ) -> None: self.vocab_size = vocab_size self.max_seq_len = max_seq_len self.pe_type = pe_type self.hidden_dim = hidden_dim self.n_heads_sa = n_heads_sa self.n_heads_ra = n_heads_ra self.n_layers = n_layers self.dropout = dropout self.dff_factor = dff_factor self.ffn_hidden_dim_mode = ffn_hidden_dim_mode self.ffn_activation = ffn_activation self.rope_theta = rope_theta self.max_rel_pos = max_rel_pos self.init_range = init_range self.init_scheme = init_scheme self.norm_type = norm_type self.norm_first = norm_first self.use_bias_qkv = use_bias_qkv self.use_bias_out = use_bias_out self.use_bias_ffn = use_bias_ffn self.tie_lm_head = tie_lm_head self.symbol_dim = symbol_dim self.n_symbols = n_symbols self.symbolic_attn_n_heads = symbolic_attn_n_heads self.symbol_retrieval = symbol_retrieval self.symbolic_use_bias = symbolic_use_bias self.shared_symbol_retriever = shared_symbol_retriever self.share_attn_params = share_attn_params self.positional_symbols_sinusoidal = positional_symbols_sinusoidal self.relative_symbols_rope = relative_symbols_rope self.relsymbolic_rel_n_heads = relsymbolic_rel_n_heads self.relsymbolic_symbolic_attn_n_heads = relsymbolic_symbolic_attn_n_heads self.relsymbolic_neighborhood_size = relsymbolic_neighborhood_size self.relsymbolic_include_self = relsymbolic_include_self self.relsymbolic_normalize_rels = relsymbolic_normalize_rels self.relsymbolic_trainable_symbols = relsymbolic_trainable_symbols self.relsymbolic_dropout = relsymbolic_dropout self.relsymbolic_rel_scale = relsymbolic_rel_scale self.relsymbolic_symbolic_attn_scale = relsymbolic_symbolic_attn_scale self.relsymbolic_use_bias = relsymbolic_use_bias self.ra_type = ra_type self.ra_n_relations = ra_n_relations self.ra_rel_activation = ra_rel_activation self.ra_symmetric_rels = ra_symmetric_rels self.sequence_boundary_policy = sequence_boundary_policy self.segment_boundary_token_id = segment_boundary_token_id self.mlm_head_enabled = mlm_head_enabled # HF base also exposes max_position_embeddings for generation utilities. self.max_position_embeddings = max_seq_len # tie_lm_head is the source of truth for weight tying; keep HF's # tie_word_embeddings in lockstep so from_pretrained never re-ties an # untied head (or fails to tie a tied one). Drop any value coming in via # kwargs (e.g. a serialized config.json) so the two cannot disagree. # NOTE: tie_lm_head controls only lm_head <-> token_embeddings tying. # The mlm_head.linear_out <-> lm_head tying is handled separately by # the get_expanded_tied_weights_keys override in modeling_dat.py. kwargs.pop("tie_word_embeddings", None) super().__init__( pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_lm_head, **kwargs, )