ctranslate2 4.7.0__cp314-cp314-macosx_11_0_arm64.whl

ctranslate2/specs/transformer_spec.py

@@ -0,0 +1,797 @@
+"""Declares specification of the Transformer model."""
+
+from typing import Optional, Tuple, Union
+
+import numpy as np
+
+from ctranslate2.specs import attention_spec, common_spec, model_spec
+
+
+class TransformerEncoderSpec(model_spec.LayerSpec):
+    def __init__(
+        self,
+        num_layers: int,
+        num_heads: int,
+        pre_norm: bool = True,
+        no_final_norm: bool = False,
+        activation: common_spec.Activation = common_spec.Activation.RELU,
+        num_source_embeddings: int = 1,
+        embeddings_merge: common_spec.EmbeddingsMerge = common_spec.EmbeddingsMerge.CONCAT,
+        layernorm_embedding: bool = False,
+        relative_position: bool = False,
+        relative_attention_bias: bool = False,
+        ffn_glu: bool = False,
+        rms_norm: bool = False,
+        multi_query_attention: bool = False,
+        num_heads_kv: Optional[int] = None,
+        head_dim: Optional[int] = None,
+        rotary_dim: Optional[int] = None,
+        rotary_interleave: bool = True,
+        rotary_scaling_type: Optional[attention_spec.RotaryScalingType] = None,
+        rotary_scaling_factor: float = 1,
+        rotary_base: float = 10000,
+        sliding_window: Optional[int] = None,
+        qk_norm: Optional[bool] = False,
+        pre_post_layer_norm: bool = False,
+    ):
+        """Initializes a Transformer encoder specification.
+
+        Args:
+          num_layers: Number of layers.
+          num_heads: Number of attention heads.
+          pre_norm: Enable the pre-norm Transformer architecture.
+          no_final_norm: Disable the final layer norm in the pre-norm architecture.
+          activation: Activation to apply in the feed-forward network.
+          num_source_embeddings: Number of source embeddings.
+          embeddings_merge: When :obj:`num_source_embeddings` > 1, specify how the
+            embeddings are merged.
+          layernorm_embedding: Apply layer normalization after the embedding layer.
+          relative_position: Use relative position representations in the self-attention
+            layers as described in https://arxiv.org/abs/1803.02155.
+          relative_attention_bias: Use relative attention bias in the self-attention
+            layers as described in the T5 paper https://arxiv.org/abs/1910.10683.
+          ffn_glu: Use gated linear units in the FFN layers as described in
+            https://arxiv.org/abs/2002.05202.
+          rms_norm: Use the root mean square layer normalization.
+          multi_query_attention: Use multi-query attention (alias for num_heads_kv=1).
+          num_heads_kv: Number of attention heads for the key and value.
+          head_dim: Number of dimensions per attention head.
+          rotary_dim: Apply rotary embeddings to these first N dimensions. If 0, rotary
+            embeddings are applied to all dimensions.
+          rotary_interleave: Interleave the head dimensions when rotary embeddings are applied.
+            Otherwise the head dimensions are sliced in half.
+          rotary_scaling_type: Type of RoPE scaling.
+          rotary_scaling_factor: Factor used in the RoPE scaling.
+          rotary_base: The base period of the rotary embeddings.
+          sliding_window: Max sequence length to retain in KV Cache.
+          qk_norm: Apply layer normalization to the query and key projections.
+          pre_post_layer_norm: Add post layer norm for each pre norm layer.
+        """
+
+        if multi_query_attention:
+            if num_heads_kv is not None and num_heads_kv != 1:
+                raise ValueError(
+                    "Enabling multi_query_attention implies num_heads_kv=1"
+                )
+            num_heads_kv = 1
+
+        self.multi_query_attention = multi_query_attention
+        self.num_heads = np.dtype("int16").type(num_heads)
+        self.pre_norm = pre_norm
+        self.activation = np.dtype("int8").type(activation)
+        self.embeddings_merge = np.dtype("int8").type(embeddings_merge)
+        self.embeddings = [
+            common_spec.EmbeddingsSpec() for _ in range(num_source_embeddings)
+        ]
+        self.scale_embeddings = True
+        if not relative_position and not relative_attention_bias:
+            self.position_encodings = PositionEncoderSpec()
+        if pre_norm and not no_final_norm:
+            self.layer_norm = common_spec.LayerNormSpec(rms_norm=rms_norm)
+        if layernorm_embedding:
+            self.layernorm_embedding = common_spec.LayerNormSpec(rms_norm=rms_norm)
+        if sliding_window is not None:
+            self.sliding_window = np.dtype("int32").type(sliding_window)
+
+        self.layer = [
+            TransformerEncoderLayerSpec(
+                relative_position=relative_position,
+                relative_attention_bias=relative_attention_bias,
+                ffn_glu=ffn_glu,
+                rms_norm=rms_norm,
+                num_heads_kv=num_heads_kv,
+                head_dim=head_dim,
+                rotary_dim=rotary_dim,
+                rotary_interleave=rotary_interleave,
+                rotary_scaling_type=rotary_scaling_type,
+                rotary_scaling_factor=rotary_scaling_factor,
+                rotary_base=rotary_base,
+                qk_norm=qk_norm,
+                pre_post_layer_norm=pre_post_layer_norm,
+            )
+            for _ in range(num_layers)
+        ]
+
+
+class TransformerDecoderSpec(model_spec.LayerSpec):
+    def __init__(
+        self,
+        num_layers: int,
+        num_heads: int,
+        pre_norm: bool = True,
+        activation: common_spec.Activation = common_spec.Activation.RELU,
+        layernorm_embedding: bool = False,
+        with_encoder_attention: bool = True,
+        no_final_norm: bool = False,
+        project_in_out: bool = False,
+        relative_position: bool = False,
+        relative_attention_bias: bool = False,
+        alignment_layer: int = -1,
+        alignment_heads: int = 1,
+        ffn_glu: bool = False,
+        rms_norm: bool = False,
+        alibi: bool = False,
+        alibi_use_positive_positions: bool = False,
+        scale_alibi: bool = False,
+        rotary_dim: Optional[int] = None,
+        rotary_interleave: bool = True,
+        rotary_scaling_type: Optional[attention_spec.RotaryScalingType] = None,
+        rotary_scaling_factor: float = 1,
+        rotary_base: float = 10000,
+        original_max_position_embeddings: int = 0,
+        max_position_embeddings: int = 0,
+        parallel_residual: bool = False,
+        shared_layer_norm: bool = False,
+        pre_post_layer_norm: bool = False,
+        multi_query_attention: bool = False,
+        num_heads_kv: Optional[int] = None,
+        head_dim: Optional[int] = None,
+        sliding_window: Optional[int] = None,
+        quant_type: Optional[common_spec.Quantization] = None,
+        quant_group_size: Optional[int] = None,
+        quant_bits: Optional[int] = None,
+        qk_norm: bool = False,
+        external_pre_post_encoder_layers: Optional[bool] = False,
+    ):
+        """Initializes a Transformer decoder specification.
+
+        Args:
+          num_layers: Number of layers.
+          num_heads: Number of attention heads.
+          pre_norm: Enable the pre-norm Transformer architecture.
+          activation: Activation to apply in the feed-forward network.
+          layernorm_embedding: Apply layer normalization after the embedding layer.
+          with_encoder_attention: Enable the encoder attention sublayers.
+          no_final_norm: Disable the final layer norm in the pre-norm architecture.
+          project_in_out: Add linear transformations after the embedding layer and before
+            the final layer.
+          relative_position: Use relative position representations in the self-attention
+            layers as described in https://arxiv.org/abs/1803.02155.
+          relative_attention_bias: Use relative attention bias in the self-attention
+            layers as described in the T5 paper https://arxiv.org/abs/1910.10683.
+          alignment_layer: Layer index selected for alignment.
+          alignment_heads: Number of attention heads selected for alignment.
+          ffn_glu: Use gated linear units in the FFN layers as described in
+            https://arxiv.org/abs/2002.05202.
+          rms_norm: Use the root mean square layer normalization.
+          alibi: Use attention with linear biases.
+          alibi_use_positive_positions: Use positive positions in the ALiBi definition.
+          scale_alibi: Apply the dot product scale factor to ALiBi.
+          rotary_dim: Apply rotary embeddings to these first N dimensions. If 0, rotary
+            embeddings are applied to all dimensions.
+          rotary_interleave: Interleave the head dimensions when rotary embeddings are applied.
+            Otherwise the head dimensions are sliced in half.
+          rotary_scaling_type: Type of RoPE scaling.
+          rotary_scaling_factor: Factor used in the RoPE scaling.
+          rotary_base: The base period of the rotary embeddings.
+          original_max_position_embeddings: The original max position embeddings
+            for Su rope embeddings
+          max_position_embeddings: The max position embeddings for Su rope embeddings
+          parallel_residual: Use parallel residual connections in each layer block, as used
+            by the GPT-J and GPT-NeoX models.
+          shared_layer_norm: When using parallel residual, share the input and post
+            attention layer norms.
+          pre_post_layer_norm: Add post layer norm for each pre norm layer
+          multi_query_attention: Use multi-query attention (alias for num_heads_kv=1).
+          num_heads_kv: Number of attention heads for the key and value.
+          sliding_window: Max sequence length to retain in KV Cache.
+          quant_type: quantization type used (like awq... for lower bit quantization)
+          quant_group_size: group size of the lower bit quantization
+          quant_bits: number of bit of the quantization (ex: 4bit)
+          external_pre_post_encoder_layers: if the encoder attention pre and processing
+            is done outside the attention.
+        """
+
+        self._config = dict()
+        if parallel_residual:
+            if not pre_norm:
+                raise ValueError("The GPT-J block expects a pre-norm architecture")
+            if with_encoder_attention:
+                raise ValueError("The GPT-J block does not have cross attention")
+
+        if multi_query_attention:
+            if num_heads_kv is not None and num_heads_kv != 1:
+                raise ValueError(
+                    "Enabling multi_query_attention implies num_heads_kv=1"
+                )
+            num_heads_kv = 1
+
+        self.num_heads = np.dtype("int16").type(num_heads)
+        self.pre_norm = pre_norm
+        self.activation = np.dtype("int8").type(activation)
+        self.alignment_layer = np.dtype("int16").type(alignment_layer)
+        self.alignment_heads = np.dtype("int16").type(alignment_heads)
+        self.embeddings = common_spec.EmbeddingsSpec()
+        self.scale_embeddings = True
+        self.scale_outputs = model_spec.OPTIONAL
+        self.alibi = alibi
+        self.alibi_use_positive_positions = alibi_use_positive_positions
+        self.scale_alibi = scale_alibi
+        if sliding_window is not None:
+            self.sliding_window = np.dtype("int32").type(sliding_window)
+        if (
+            not relative_position
+            and not relative_attention_bias
+            and not alibi
+            and rotary_dim is None
+        ):
+            self.position_encodings = PositionEncoderSpec()
+        if pre_norm and not no_final_norm:
+            self.layer_norm = common_spec.LayerNormSpec(rms_norm=rms_norm)
+        if layernorm_embedding:
+            self.layernorm_embedding = common_spec.LayerNormSpec(rms_norm=rms_norm)
+        self.projection = common_spec.LinearSpec()
+        self.layer = [
+            TransformerDecoderLayerSpec(
+                with_encoder_attention=with_encoder_attention,
+                relative_position=relative_position,
+                relative_attention_bias=relative_attention_bias,
+                ffn_glu=ffn_glu,
+                rms_norm=rms_norm,
+                rotary_dim=rotary_dim,
+                rotary_interleave=rotary_interleave,
+                rotary_scaling_type=rotary_scaling_type,
+                rotary_scaling_factor=rotary_scaling_factor,
+                rotary_base=rotary_base,
+                original_max_position_embeddings=original_max_position_embeddings,
+                max_position_embeddings=max_position_embeddings,
+                parallel_residual=parallel_residual,
+                shared_layer_norm=shared_layer_norm,
+                pre_post_layer_norm=pre_post_layer_norm,
+                num_heads_kv=num_heads_kv,
+                head_dim=head_dim,
+                sliding_window=sliding_window,
+                qk_norm=qk_norm,
+                external_pre_post_encoder_layers=external_pre_post_encoder_layers,
+            )
+            for _ in range(num_layers)
+        ]
+        self.start_from_zero_embedding = False
+        self._config["multi_query_attention"] = multi_query_attention or (
+            num_heads_kv != num_heads
+        )
+
+        if project_in_out:
+            self.project_in = common_spec.LinearSpec()
+            self.project_out = common_spec.LinearSpec()
+
+        if quant_type:
+            self._config["quantization_type"] = quant_type
+            self._config["quantization_bits"] = quant_bits
+            self._config["quantization_group_size"] = quant_group_size
+
+    @property
+    def config(self):
+        return self._config
+
+
+class TransformerEncoderLayerSpec(model_spec.LayerSpec):
+    def __init__(
+        self,
+        relative_position=False,
+        relative_attention_bias=False,
+        ffn_glu=False,
+        rms_norm=False,
+        num_heads_kv=None,
+        head_dim=None,
+        sliding_window=None,
+        rotary_dim: Optional[int] = None,
+        rotary_interleave: bool = True,
+        rotary_scaling_type: Optional[attention_spec.RotaryScalingType] = None,
+        rotary_scaling_factor: float = 1,
+        rotary_base: float = 10000,
+        qk_norm=False,
+        pre_post_layer_norm: bool = False,
+    ):
+        self.self_attention = attention_spec.MultiHeadAttentionSpec(
+            self_attention=True,
+            relative_position=relative_position,
+            relative_attention_bias=relative_attention_bias,
+            rms_norm=rms_norm,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            sliding_window=sliding_window,
+            rotary_dim=rotary_dim,
+            rotary_interleave=rotary_interleave,
+            rotary_scaling_type=rotary_scaling_type,
+            rotary_scaling_factor=rotary_scaling_factor,
+            rotary_base=rotary_base,
+            qk_norm=qk_norm,
+        )
+        self.ffn = FeedForwardSpec(glu=ffn_glu, rms_norm=rms_norm)
+
+        if pre_post_layer_norm:
+            self.input_layer_norm = common_spec.LayerNormSpec(rms_norm=rms_norm)
+            self.post_attention_layer_norm = common_spec.LayerNormSpec(
+                rms_norm=rms_norm
+            )
+            self.pre_feedforward_layer_norm = common_spec.LayerNormSpec(
+                rms_norm=rms_norm
+            )
+            self.post_feedforward_layer_norm = common_spec.LayerNormSpec(
+                rms_norm=rms_norm
+            )
+
+            delattr(self.self_attention, "layer_norm")
+            delattr(self.ffn, "layer_norm")
+
+
+class TransformerDecoderLayerSpec(model_spec.LayerSpec):
+    def __init__(
+        self,
+        with_encoder_attention=True,
+        relative_position=False,
+        relative_attention_bias=False,
+        ffn_glu=False,
+        rms_norm=False,
+        rotary_dim=None,
+        rotary_interleave=True,
+        rotary_scaling_type=None,
+        rotary_scaling_factor=1,
+        rotary_base=10000,
+        original_max_position_embeddings=0,
+        max_position_embeddings=0,
+        parallel_residual=False,
+        shared_layer_norm=False,
+        pre_post_layer_norm=False,
+        num_heads_kv=None,
+        head_dim=None,
+        sliding_window=None,
+        qk_norm=False,
+        external_pre_post_encoder_layers=False,
+    ):
+        self.self_attention = attention_spec.MultiHeadAttentionSpec(
+            self_attention=True,
+            relative_position=relative_position,
+            relative_attention_bias=relative_attention_bias,
+            rms_norm=rms_norm,
+            rotary_dim=rotary_dim,
+            rotary_interleave=rotary_interleave,
+            rotary_scaling_type=rotary_scaling_type,
+            rotary_scaling_factor=rotary_scaling_factor,
+            rotary_base=rotary_base,
+            original_max_position_embeddings=original_max_position_embeddings,
+            max_position_embeddings=max_position_embeddings,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            sliding_window=sliding_window,
+            qk_norm=qk_norm,
+        )
+
+        if with_encoder_attention:
+            self.attention = attention_spec.MultiHeadAttentionSpec(
+                rms_norm=rms_norm,
+                num_heads_kv=num_heads_kv,
+                head_dim=head_dim,
+                sliding_window=sliding_window,
+                qk_norm=qk_norm,
+                has_norm=external_pre_post_encoder_layers is False,
+            )
+
+        self.ffn = FeedForwardSpec(glu=ffn_glu, rms_norm=rms_norm)
+
+        if parallel_residual:
+            if shared_layer_norm:
+                self.shared_layer_norm = common_spec.LayerNormSpec()
+            else:
+                self.input_layer_norm = common_spec.LayerNormSpec()
+                self.post_attention_layer_norm = common_spec.LayerNormSpec()
+
+            delattr(self.self_attention, "layer_norm")
+            delattr(self.ffn, "layer_norm")
+
+        if pre_post_layer_norm:
+            # Self-attention layer norms
+            self.input_layer_norm = common_spec.LayerNormSpec(rms_norm=rms_norm)
+            self.post_attention_layer_norm = common_spec.LayerNormSpec(
+                rms_norm=rms_norm
+            )
+
+            if with_encoder_attention and external_pre_post_encoder_layers:
+                self.external_post_encoder_attention_layer_norm = (
+                    common_spec.LayerNormSpec(rms_norm=rms_norm)
+                )
+                self.external_pre_encoder_attention_layer_norm = (
+                    common_spec.LayerNormSpec(rms_norm=rms_norm)
+                )
+
+            # Feed-forward layer norms
+            self.pre_feedforward_layer_norm = common_spec.LayerNormSpec(
+                rms_norm=rms_norm
+            )
+            self.post_feedforward_layer_norm = common_spec.LayerNormSpec(
+                rms_norm=rms_norm
+            )
+
+            delattr(self.self_attention, "layer_norm")
+            delattr(self.ffn, "layer_norm")
+
+
+class FeedForwardSpec(model_spec.LayerSpec):
+    def __init__(self, glu=False, rms_norm=False):
+        self.layer_norm = common_spec.LayerNormSpec(rms_norm=rms_norm)
+        self.linear_0 = common_spec.LinearSpec()
+        self.linear_1 = common_spec.LinearSpec()
+        if glu:
+            self.linear_0_noact = common_spec.LinearSpec()
+
+
+class PositionEncoderSpec(model_spec.LayerSpec):
+    def __init__(self):
+        self.encodings = model_spec.OPTIONAL
+
+
+class TransformerConfig(model_spec.SequenceToSequenceModelConfig):
+    """Configuration for Transformer models."""
+
+    def __init__(self, layer_norm_epsilon: Optional[float] = None, **kwargs):
+        """Initializes the configuration for Transformer models.
+
+        Args:
+          layer_norm_epsilon: The layer norm epsilon value.
+          **kwargs: Additional configuration.
+        """
+        super().__init__(layer_norm_epsilon=layer_norm_epsilon, **kwargs)
+
+
+class TransformerSpec(model_spec.SequenceToSequenceModelSpec):
+    """Describes a Transformer model.
+
+    The specification is invariant to hidden dimensions but requires to
+    explicitly set the number of layers and attention heads.
+    """
+
+    def __init__(
+        self, encoder: TransformerEncoderSpec, decoder: TransformerDecoderSpec
+    ):
+        """Initializes a Transformer model specification.
+
+        Args:
+          encoder: The encoder specification.
+          decoder: The decoder specification.
+        """
+        if not isinstance(encoder, TransformerEncoderSpec):
+            raise TypeError("encoder argument must be a TransformerEncoderSpec")
+        if not isinstance(decoder, TransformerDecoderSpec):
+            raise TypeError("decoder argument must be a TransformerDecoderSpec")
+
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self._config.add_attribute(
+            "multi_query_attention", self.encoder.multi_query_attention
+        )
+
+    @classmethod
+    def from_config(
+        cls,
+        num_layers: Union[int, Tuple[int, int]],
+        num_heads: int,
+        with_relative_position: bool = False,
+        pre_norm: bool = True,
+        no_final_norm: bool = False,
+        activation: common_spec.Activation = common_spec.Activation.RELU,
+        alignment_layer: int = -1,
+        alignment_heads: int = 1,
+        num_source_embeddings: int = 1,
+        embeddings_merge: common_spec.EmbeddingsMerge = common_spec.EmbeddingsMerge.CONCAT,
+        layernorm_embedding: bool = False,
+        relative_attention_bias: bool = False,
+        ffn_glu: bool = False,
+        rms_norm: bool = False,
+        multi_query_attention: bool = False,
+    ):
+        """Creates a Transformer model specification.
+
+        Args:
+          num_layers: Number of encoder and decoder layers, or a 2-tuple if the
+            number is different.
+          num_heads: Number of attention heads.
+          with_relative_position: Use relative position representations in the self-attention
+            layers as described in https://arxiv.org/abs/1803.02155.
+          pre_norm: Enable the pre-norm Transformer architecture.
+          no_final_norm: Disable the final layer norm in the pre-norm architecture.
+          activation: Activation to apply in the feed-forward network.
+          alignment_layer: Layer index selected for alignment.
+          alignment_heads: Number of attention heads selected for alignment.
+          num_source_embeddings: Number of source embeddings.
+          embeddings_merge: When :obj:`num_source_embeddings` > 1, specify how the
+            embeddings are merged.
+          layernorm_embedding: Apply layer normalization after the embedding layer.
+          relative_attention_bias: Use relative attention bias in the self-attention
+            layers as described in the T5 paper https://arxiv.org/abs/1910.10683.
+          ffn_glu: Use gated linear units in the FFN layer as described in
+            https://arxiv.org/abs/2002.05202.
+          rms_norm: Use the root mean square layer normalization.
+          multi_query_attention: Use multi-query attention.
+        """
+        if isinstance(num_layers, (list, tuple)):
+            num_encoder_layers, num_decoder_layers = num_layers
+        else:
+            num_encoder_layers, num_decoder_layers = num_layers, num_layers
+
+        encoder = TransformerEncoderSpec(
+            num_encoder_layers,
+            num_heads,
+            pre_norm=pre_norm,
+            no_final_norm=no_final_norm,
+            activation=activation,
+            num_source_embeddings=num_source_embeddings,
+            embeddings_merge=embeddings_merge,
+            layernorm_embedding=layernorm_embedding,
+            relative_position=with_relative_position,
+            relative_attention_bias=relative_attention_bias,
+            ffn_glu=ffn_glu,
+            rms_norm=rms_norm,
+            multi_query_attention=multi_query_attention,
+        )
+
+        decoder = TransformerDecoderSpec(
+            num_decoder_layers,
+            num_heads,
+            pre_norm=pre_norm,
+            no_final_norm=no_final_norm,
+            activation=activation,
+            layernorm_embedding=layernorm_embedding,
+            relative_position=with_relative_position,
+            relative_attention_bias=relative_attention_bias,
+            alignment_layer=alignment_layer,
+            alignment_heads=alignment_heads,
+            ffn_glu=ffn_glu,
+            rms_norm=rms_norm,
+            multi_query_attention=multi_query_attention,
+        )
+
+        return cls(encoder, decoder)
+
+    @property
+    def name(self):
+        return "TransformerSpec"
+
+    @property
+    def revision(self):
+        return 7
+
+    def get_default_config(self):
+        return TransformerConfig()
+
+    def get_source_vocabulary_size(self):
+        return [spec.weight.shape[0] for spec in self.encoder.embeddings]
+
+    def get_target_vocabulary_size(self):
+        return self.decoder.embeddings.weight.shape[0]
+
+
+class TransformerDecoderModelConfig(model_spec.LanguageModelConfig):
+    """Configuration for Transformer decoder models."""
+
+    def __init__(self, layer_norm_epsilon: Optional[float] = None, **kwargs):
+        """Initializes the configuration for Transformer decoder models.
+
+        Args:
+          layer_norm_epsilon: The layer norm epsilon value.
+          **kwargs: Additional configuration.
+        """
+        super().__init__(layer_norm_epsilon=layer_norm_epsilon, **kwargs)
+
+
+class TransformerDecoderModelSpec(model_spec.LanguageModelSpec):
+    """Describes a Transformer decoder model (e.g. GPT-2)."""
+
+    def __init__(self, decoder: TransformerDecoderSpec):
+        """Initializes a Transformer decoder model specification.
+
+        Args:
+          decoder: The decoder specification.
+        """
+        if not isinstance(decoder, TransformerDecoderSpec):
+            raise TypeError("decoder argument must be a TransformerDecoderSpec")
+
+        super().__init__()
+        self.decoder = decoder
+        for key, value in self.decoder.config.items():
+            self._config.add_attribute(key, value)
+
+    @classmethod
+    def from_config(
+        cls,
+        num_layers: int,
+        num_heads: int,
+        pre_norm: bool = True,
+        activation: common_spec.Activation = common_spec.Activation.RELU,
+        layernorm_embedding: bool = False,
+        no_final_norm: bool = False,
+        project_in_out: bool = False,
+        with_relative_position: bool = False,
+        ffn_glu: bool = False,
+        rms_norm: bool = False,
+        alibi: bool = False,
+        alibi_use_positive_positions: bool = False,
+        scale_alibi: bool = False,
+        rotary_dim: Optional[int] = None,
+        rotary_interleave: bool = True,
+        rotary_scaling_type: Optional[attention_spec.RotaryScalingType] = None,
+        rotary_scaling_factor: float = 1,
+        rotary_base: float = 10000,
+        original_max_position_embeddings: int = 0,
+        max_position_embeddings: int = 0,
+        parallel_residual: bool = False,
+        shared_layer_norm: bool = False,
+        pre_post_layer_norm: bool = False,
+        multi_query_attention: bool = False,
+        num_heads_kv: Optional[int] = None,
+        head_dim: Optional[int] = None,
+        sliding_window: Optional[int] = None,
+        quant_type: Optional[common_spec.Quantization] = None,
+        quant_group_size: Optional[int] = None,
+        quant_bits: Optional[int] = None,
+        qk_norm: bool = False,
+    ):
+        """Creates a Transformer decoder model specification.
+
+        Args:
+          num_layers: Number of decoder layers.
+          num_heads: Number of attention heads.
+          pre_norm: Enable the pre-norm Transformer architecture.
+          activation: Activation to apply in the feed-forward network.
+          layernorm_embedding: Apply layer normalization after the embedding layer.
+          no_final_norm: Do not apply layer normalization after the last decoder block.
+          project_in_out: Add a linear layer after the embedding layer and another one
+            before the final output projection.
+          with_relative_position: Enable relative position representations modules.
+          ffn_glu: Use gated linear units in the FFN layers as described in
+            https://arxiv.org/abs/2002.05202.
+          rms_norm: Use the root mean square layer normalization.
+          alibi: Use attention with linear biases.
+          alibi_use_positive_positions: Use positive positions in the ALiBi definition.
+          scale_alibi: Apply the dot product scale factor to ALiBi.
+          rotary_dim: Apply rotary embeddings to these first N dimensions. If 0, rotary
+            embeddings are applied to all dimensions.
+          rotary_interleave: Interleave the head dimensions when rotary embeddings are applied.
+            Otherwise the head dimensions are sliced in half.
+          rotary_scaling_type: Type of RoPE scaling.
+          rotary_scaling_factor: Factor used in the RoPE scaling.
+          rotary_base: The base period of the rotary embeddings.
+          original_max_position_embeddings: The original max position embeddings
+            for Su rope embeddings
+          max_position_embeddings: The max position embeddings for Su rope embeddings
+          parallel_residual: Use parallel residual connections in each layer block, as used
+            by the GPT-J and GPT-NeoX models.
+          shared_layer_norm: When using parallel residual, share the input and post
+            attention layer norms.
+          pre_post_layer_norm: add post layer norm for each pre norm layer
+          multi_query_attention: Use multi-query attention (alias for num_heads_kv=1).
+          num_heads_kv: Number of attention heads for the key and value.
+          head_dim: Number of head
+          sliding_window: max sequence length to retain KV cache
+          quant_type: quantization type used (like awq... for lower bit quantization)
+          quant_group_size: group size of the lower bit quantization
+          quant_bits: number of bit of the quantization (ex: 4bit)
+        """
+        decoder = TransformerDecoderSpec(
+            num_layers,
+            num_heads,
+            pre_norm=pre_norm,
+            activation=activation,
+            layernorm_embedding=layernorm_embedding,
+            with_encoder_attention=False,
+            no_final_norm=no_final_norm,
+            project_in_out=project_in_out,
+            relative_position=with_relative_position,
+            ffn_glu=ffn_glu,
+            rms_norm=rms_norm,
+            alibi=alibi,
+            alibi_use_positive_positions=alibi_use_positive_positions,
+            scale_alibi=scale_alibi,
+            rotary_dim=rotary_dim,
+            rotary_interleave=rotary_interleave,
+            rotary_scaling_type=rotary_scaling_type,
+            rotary_scaling_factor=rotary_scaling_factor,
+            rotary_base=rotary_base,
+            original_max_position_embeddings=original_max_position_embeddings,
+            max_position_embeddings=max_position_embeddings,
+            parallel_residual=parallel_residual,
+            shared_layer_norm=shared_layer_norm,
+            pre_post_layer_norm=pre_post_layer_norm,
+            multi_query_attention=multi_query_attention,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            sliding_window=sliding_window,
+            quant_type=quant_type,
+            quant_group_size=quant_group_size,
+            quant_bits=quant_bits,
+            qk_norm=qk_norm,
+        )
+
+        return cls(decoder)
+
+    @property
+    def name(self):
+        return "TransformerDecoderSpec"
+
+    @property
+    def revision(self):
+        return 8
+
+    def get_default_config(self):
+        return TransformerDecoderModelConfig()
+
+    def get_vocabulary_size(self):
+        return self.decoder.embeddings.weight.shape[0]
+
+
+class TransformerEncoderModelConfig(model_spec.LanguageModelConfig):
+    """Configuration for Transformer encoder models."""
+
+    def __init__(self, layer_norm_epsilon: Optional[float] = None, **kwargs):
+        """Initializes the configuration for Transformer encoder models.
+
+        Args:
+          layer_norm_epsilon: The layer norm epsilon value.
+          **kwargs: Additional configuration.
+        """
+        super().__init__(layer_norm_epsilon=layer_norm_epsilon, **kwargs)
+
+
+class TransformerEncoderModelSpec(model_spec.LanguageModelSpec):
+    """Describes a Transformer encoder model (e.g. BERT)."""
+
+    def __init__(
+        self,
+        encoder: TransformerEncoderSpec,
+        pooling_layer: bool = False,
+        pooling_activation: common_spec.Activation = common_spec.Activation.Tanh,
+    ):
+        """Initializes a Transformer encoder model specification.
+
+        Args:
+          encoder: The encoder specification.
+          pooling_layer: Add the pooling layer.
+          pooling_activation: The activation to apply after the pooling layer.
+        """
+        if not isinstance(encoder, TransformerEncoderSpec):
+            raise TypeError("encoder argument must be a TransformerEncoderSpec")
+
+        super().__init__()
+        self.encoder = encoder
+        self._config.add_attribute(
+            "multi_query_attention", self.encoder.multi_query_attention
+        )
+
+        if pooling_layer:
+            self.pooler_dense = common_spec.LinearSpec()
+            self.pooler_activation = np.dtype("int8").type(pooling_activation)
+
+    @property
+    def name(self):
+        return "TransformerEncoderSpec"
+
+    @property
+    def revision(self):
+        return 1
+
+    def get_default_config(self):
+        return TransformerEncoderModelConfig()
+
+    def get_vocabulary_size(self):
+        return self.encoder.embeddings[0].weight.shape[0]