configuration_aliceai_t5.py

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