from copy import deepcopy from typing import Any from transformers.configuration_utils import PretrainedConfig class AliceAIT5ModuleConfig(PretrainedConfig): """Attention and RMS normalization settings for one stack.""" model_type = "aliceai_t5_module" keys_to_ignore_at_inference = ["past_key_values"] def __init__( self, vocab_size=134016, hidden_size=3072, num_hidden_layers=24, num_attention_heads=24, num_key_value_heads=24, head_dim=128, max_position_embeddings=20005, initializer_range=0.02, norm_eps=9.999999747378752e-06, use_cache=True, pad_token_id=1, eos_token_id=2, bos_token_id=2, tie_word_embeddings=True, rope_theta=10000.0, rope_parameters=None, rope_scaling=None, attention_bias=False, attention_dropout=0.0, query_pre_attn_scalar=256, sliding_window=4096, layer_types=None, final_logit_softcapping=None, attn_logit_softcapping=None, fp32_residual=False, _attn_implementation="eager", **kwargs, ): if not kwargs.pop("use_rms_norm", True) or kwargs.pop("layer_norm_bias", False): raise ValueError("AliceAIT5 only supports RMSNorm without a bias.") for legacy_field in ("intermediate_size", "hidden_activation", "mlp_bias"): kwargs.pop(legacy_field, None) # RoPE validation in PretrainedConfig needs the architecture fields first. self.vocab_size = vocab_size self.max_position_embeddings = max_position_embeddings self.hidden_size = hidden_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.head_dim = head_dim self.num_key_value_heads = num_key_value_heads self.initializer_range = initializer_range self.norm_eps = norm_eps self.use_cache = use_cache self.rope_theta = rope_theta self.rope_parameters = deepcopy(rope_parameters if rope_parameters is not None else rope_scaling) self.attention_bias = attention_bias self.attention_dropout = attention_dropout self.query_pre_attn_scalar = query_pre_attn_scalar self.sliding_window = sliding_window self.final_logit_softcapping = final_logit_softcapping self.attn_logit_softcapping = attn_logit_softcapping self.layer_types = ( list(layer_types) if layer_types is not None else ["full_attention"] * self.num_hidden_layers ) self.fp32_residual = fp32_residual self._validate_architecture() attn_implementation = kwargs.pop("attn_implementation", _attn_implementation) super().__init__( pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, attn_implementation=attn_implementation, **kwargs, ) def _validate_architecture(self): if len(self.layer_types) < self.num_hidden_layers: raise ValueError("layer_types must contain at least one entry per hidden layer") # Ignore layer entries beyond the configured depth. self.layer_types = self.layer_types[: self.num_hidden_layers] invalid_layer_types = set(self.layer_types) - {"full_attention", "sliding_attention"} if invalid_layer_types: raise ValueError(f"Unsupported attention layer types: {sorted(invalid_layer_types)}") if self.num_attention_heads % self.num_key_value_heads: raise ValueError("num_attention_heads must be divisible by num_key_value_heads") if self.num_attention_heads * self.head_dim != self.hidden_size: raise ValueError("num_attention_heads * head_dim must equal hidden_size") class AliceAIT5Config(PretrainedConfig): """Encoder-decoder configuration with optional shared embeddings.""" model_type = "aliceai_t5" keys_to_ignore_at_inference = ["past_key_values"] sub_configs = {"encoder": AliceAIT5ModuleConfig, "decoder": AliceAIT5ModuleConfig} def __init__( self, encoder: AliceAIT5ModuleConfig | dict[Any, Any] | None = None, decoder: AliceAIT5ModuleConfig | dict[Any, Any] | None = None, is_encoder_decoder: bool = True, dropout_rate: float = 0.0, attention_dropout: float = 0.0, shared_embeddings: bool = True, tie_word_embeddings: bool = True, _attn_implementation: str = "eager", **kwargs, ): if isinstance(encoder, dict): encoder = AliceAIT5ModuleConfig(**encoder) elif encoder is None: encoder = AliceAIT5ModuleConfig() else: assert isinstance(encoder, AliceAIT5ModuleConfig), f"{type(encoder)} is not supported." if isinstance(decoder, dict): decoder = AliceAIT5ModuleConfig(**decoder) elif decoder is None: decoder = encoder else: assert isinstance(decoder, AliceAIT5ModuleConfig), f"{type(decoder)} is not supported." # Keep encoder and decoder settings independent. encoder = deepcopy(encoder) decoder = deepcopy(decoder) resolved_use_cache = kwargs.pop("use_cache", decoder.use_cache) encoder.is_decoder = False encoder.dropout_rate = dropout_rate encoder.attention_dropout = attention_dropout self.encoder = encoder decoder.is_decoder = True decoder.use_cache = resolved_use_cache decoder.dropout_rate = dropout_rate decoder.attention_dropout = attention_dropout decoder.cross_attention_hidden_size = encoder.hidden_size self.decoder = decoder for special_token_key in ["bos_token_id", "pad_token_id", "eos_token_id"]: if special_token_key not in kwargs: kwargs[special_token_key] = getattr(decoder, special_token_key) self.is_encoder_decoder = is_encoder_decoder self.use_cache = resolved_use_cache self.initializer_range = kwargs.get("initializer_range", decoder.initializer_range) self.dropout_rate = dropout_rate self.attention_dropout = attention_dropout self.shared_embeddings = shared_embeddings decoder.has_embeddings = not shared_embeddings encoder.has_embeddings = True self.tie_word_embeddings = tie_word_embeddings self.vocab_size = encoder.vocab_size attn_implementation = kwargs.pop("attn_implementation", _attn_implementation) super().__init__( is_encoder_decoder=is_encoder_decoder, tie_word_embeddings=tie_word_embeddings, use_cache=self.use_cache, attn_implementation=attn_implementation, **kwargs, ) def get_text_config(self, decoder=None, encoder=None) -> "PretrainedConfig": if decoder and encoder: raise ValueError("Specify either decoder=True or encoder=True, not both") if decoder: return self.decoder if encoder: return self.encoder return self