pys2t_transformer.py

#!/usr/bin/env python3

import logging
import math
import torch
from functools import reduce

import torch.nn as nn
from fairseq import checkpoint_utils
from fairseq.data.data_utils import lengths_to_padding_mask
from fairseq.models import (
    FairseqEncoder,
    register_model,
    register_model_architecture,
)
from fairseq.models.speech_to_text import S2TTransformerModel

from fairseq.modules import (
    FairseqDropout,
    LayerNorm,
    PositionalEmbedding,
    PyramidTransformerEncoderLayer,
    MultiheadAttention,
)

logger = logging.getLogger(__name__)


def lengths_to_padding_mask_with_maxlen(lens, max_length):
    bsz = lens.size(0)
    mask = torch.arange(max_length).to(lens.device).view(1, max_length)
    mask = mask.expand(bsz, -1) >= lens.view(bsz, 1).expand(-1, max_length)
    return mask


class ReducedEmbed(nn.Module):
    # Reduced embedding for Pyramid Transformer
    def __init__(
        self,
        reduced_way: str,
        embed_norm: bool,
        in_channels: int,
        out_channels: int,
        kernel_sizes: int,
        stride: int,
        padding: int,
    ):
        super().__init__()

        self.stride = stride
        self.reduced_way = reduced_way
        if self.reduced_way == "conv":
            self.conv = nn.Conv1d(in_channels, out_channels, kernel_sizes, stride=stride, padding=padding)
        elif self.reduced_way == "glu":
            self.conv = nn.Conv1d(in_channels, out_channels * 2, kernel_sizes, stride=stride, padding=padding)
            self.glu = nn.GLU(dim=1)
        elif self.reduced_way == "proj":
            self.proj = nn.Linear(2 * in_channels, out_channels, bias=False)
        else:
            logger.error("Unsupported reduced way!")

        self.embed_norm = embed_norm
        if self.embed_norm:
            if self.reduced_way == "proj":
                self.norm = LayerNorm(2 * in_channels)
            else:
                self.norm = LayerNorm(out_channels)

    def forward(self, x, lengths):
        seq_len, bsz, dim = x.size()
        assert seq_len % self.stride == 0, "The sequence length %d must be a multiple of %d." % (seq_len, self.stride)

        padding_mask = lengths_to_padding_mask_with_maxlen(lengths, seq_len)    # bsz, seq_len
        mask_pad = padding_mask.unsqueeze(2)
        # mask batch padding
        if mask_pad is not None:
            x = x.transpose(0, 1)
            x.masked_fill_(mask_pad, 0.0)
            x = x.transpose(0, 1)

        if self.reduced_way == "proj":
            x = x.transpose(0, 1).contiguous().view(bsz, int(seq_len / 2), -1)
            x = self.proj(self.norm(x))
            x = x.transpose(0, 1)
        else:
            x = x.permute(1, 2, 0)  # B * D * T
            x = self.conv(x)
            if self.reduced_way == "glu":
                x = self.glu(x)
            x = x.permute(2, 0, 1)  # T * B * D
        if self.embed_norm:
            x = self.norm(x)

        lengths = lengths / self.stride

        padding_mask = lengths_to_padding_mask_with_maxlen(lengths, x.size(0))
        mask_pad = padding_mask.unsqueeze(2)
        # mask batch padding
        if mask_pad is not None:
            x = x.transpose(0, 1)
            x.masked_fill_(mask_pad, 0.0)
            x = x.transpose(0, 1)

        return x, lengths, padding_mask


class BlockFuse(nn.Module):

    def __init__(self, embed_dim, prev_embed_dim, dropout, num_head, fuse_way="add"):
        super().__init__()

        self.attn = MultiheadAttention(
            embed_dim,
            num_head,
            kdim=prev_embed_dim,
            vdim=prev_embed_dim,
            dropout=dropout,
            encoder_decoder_attention=True,
        )
        self.q_layer_norm = LayerNorm(embed_dim)
        self.kv_layer_norm = LayerNorm(prev_embed_dim)
        self.fuse_way = fuse_way
        if self.fuse_way == "gated":
            self.gate_linear = nn.Linear(2 * embed_dim, embed_dim)

    def forward(self, x, state, padding):
        residual = x
        x = self.q_layer_norm(x)
        state = self.kv_layer_norm(state)
        x, attn = self.attn(
            query=x,
            key=state,
            value=state,
            key_padding_mask=padding,
            static_kv=True,
        )

        if self.fuse_way == "add":
            x = residual + x
        elif self.fuse_way == "gated":
            coef = (self.gate_linear(torch.cat([x, residual], dim=-1))).sigmoid()
            x = coef * x + (1 - coef) * residual

        return x


@register_model("pys2t_transformer")
class PYS2TTransformerModel(S2TTransformerModel):
    """Adapted Transformer model (https://arxiv.org/abs/1706.03762) for
    speech-to-text tasks. The Transformer encoder/decoder remains the same.
    A trainable input subsampler is prepended to the Transformer encoder to
    project inputs into the encoder dimension as well as downsample input
    sequence for computational efficiency."""

    def __init__(self, encoder, decoder):
        super().__init__(encoder, decoder)

    @staticmethod
    def add_args(parser):
        """Add model-specific arguments to the parser."""
        S2TTransformerModel.add_args(parser)
        parser.add_argument(
            "--pyramid-stages",
            type=int,
            help="the number of the stage",
        )
        parser.add_argument(
            "--pyramid-layers",
            type=str,
            help="the number of the encoder layers",
        )
        parser.add_argument(
            "--pyramid-sr-ratios",
            type=str,
            help="the ratio of the subsampling",
        )
        parser.add_argument(
            "--pyramid-attn-sample-ratio",
            type=str,
            help="the ratio of the subsampling in the self attention module",
        )
        parser.add_argument(
            "--pyramid-reduced-embed",
            type=str,
            choices=["glu", "conv", "proj"],
            help="the reduced way of the embedding",
        )
        parser.add_argument(
            "--pyramid-embed-norm",
            action="store_true",
            help="use layer norm in reduced embedding",
        )
        parser.add_argument(
            "--pyramid-block-attn",
            action="store_true",
            help="use block attention",
        )
        parser.add_argument(
            "--pyramid-fuse-way",
            type=str,
            help="fused way for block attention",
        )
        parser.add_argument(
            "--pyramid-position-embed",
            type=str,
            help="use the position embedding or not",
        )
        parser.add_argument(
            "--pyramid-embed-dims",
            type=str,
            help="the embedding dimension",
        )
        parser.add_argument(
            "--pyramid-kernel-sizes",
            type=str,
            help="the kernel size of the reduced embedding",
        )
        parser.add_argument(
            "--pyramid-ffn-ratios",
            type=str,
            help="the ratio of the ffn",
        )
        parser.add_argument(
            "--pyramid-heads",
            type=str,
            help="the number of the attention heads",
        )
        parser.add_argument(
            "--pyramid-use-ppm",
            action="store_true",
            help="use ppm",
        )

        parser.add_argument(
            "--ctc-layer",
            type=int,
            help="the position of the ctc loss",
        )
        pass

    @classmethod
    def build_encoder(cls, args, task=None, embed_tokens=None):
        encoder = PyS2TTransformerEncoder(args, task, embed_tokens)
        if getattr(args, "load_pretrained_encoder_from", None):
            encoder = checkpoint_utils.load_pretrained_component_from_model(
                component=encoder, checkpoint=args.load_pretrained_encoder_from, strict=False
            )
            logger.info(
                f"loaded pretrained encoder from: "
                f"{args.load_pretrained_encoder_from}"
            )
        return encoder


class PyS2TTransformerEncoder(FairseqEncoder):
    """Speech-to-text Pyramid Transformer encoder"""

    def __init__(self, args, task=None, embed_tokens=None):
        super().__init__(None)

        self.padding_idx = 1
        self.attn_type = getattr(args, "encoder_attention_type", "selfattn")
        self.embed_dim = args.encoder_embed_dim

        self.dropout = FairseqDropout(
                p=args.dropout, module_name=self.__class__.__name__
        )
        self.pyramid_stages = getattr(args, "pyramid_stages", 4)

        self.pyramid_layers = [int(n) for n in args.pyramid_layers.split("_")]
        self.pyramid_sr_ratios = [int(n) for n in args.pyramid_sr_ratios.split("_")]
        self.pyramid_attn_sample_ratios = [int(n) for n in args.pyramid_attn_sample_ratios.split("_")]
        self.pyramid_embed_dims = [int(n) for n in args.pyramid_embed_dims.split("_")]
        self.pyramid_position_embed = [int(n) for n in args.pyramid_position_embed.split("_")]
        self.pyramid_kernel_sizes = [int(n) for n in args.pyramid_kernel_sizes.split("_")]
        self.pyramid_ffn_ratios = [int(n) for n in args.pyramid_ffn_ratios.split("_")]
        self.pyramid_heads = [int(n) for n in args.pyramid_heads.split("_")]
        self.pyramid_reduced_embed = args.pyramid_reduced_embed
        self.pyramid_embed_norm = args.pyramid_embed_norm

        self.pyramid_block_attn = getattr(args, "pyramid_block_attn", False)
        self.pyramid_fuse_way = getattr(args, "pyramid_fuse_way", "add")
        self.use_ppm = getattr(args, "pyramid_use_ppm", False)

        for i in range(self.pyramid_stages):
            num_layers = self.pyramid_layers[i]
            sr_ratio = self.pyramid_sr_ratios[i]
            attn_sample_ratio = self.pyramid_attn_sample_ratios[i]
            embed_dim = self.pyramid_embed_dims[i]
            kernel_size = self.pyramid_kernel_sizes[i]
            ffn_ratio = self.pyramid_ffn_ratios[i]
            num_head = self.pyramid_heads[i]
            use_pos_embed = self.pyramid_position_embed[i]

            if i == 0:
                self.embed_scale = math.sqrt(embed_dim)
                if args.no_scale_embedding:
                    self.embed_scale = 1.0

            reduced_embed = ReducedEmbed(
                self.pyramid_reduced_embed,
                self.pyramid_embed_norm,
                # self.pyramid_embed_norm if i != 0 else False,
                args.input_feat_per_channel * args.input_channels if i == 0 else self.pyramid_embed_dims[i-1],
                embed_dim,
                kernel_sizes=kernel_size,
                stride=sr_ratio,
                padding=kernel_size // 2,
            )
            if use_pos_embed:
                pos_embed = PositionalEmbedding(
                    args.max_source_positions, embed_dim,
                    self.padding_idx, pos_emb_type=self.attn_type
                )
            else:
                pos_embed = None

            block = nn.ModuleList([
                PyramidTransformerEncoderLayer(args, embed_dim, embed_dim * ffn_ratio, num_head, attn_sample_ratio)
                for _ in range(num_layers)])

            block_attn = None
            if self.pyramid_block_attn:
                if i != 0:
                    block_attn = BlockFuse(embed_dim, self.pyramid_embed_dims[i-1],
                                           args.dropout, num_head, self.pyramid_fuse_way)

            if self.use_ppm:
                ppm_layer_norm = LayerNorm(embed_dim)
                ppm = nn.Sequential(
                    nn.Conv1d(embed_dim, self.embed_dim, kernel_size=1, bias=False),
                    nn.BatchNorm1d(self.embed_dim),
                    nn.ReLU(),
                )
            else:
                ppm_layer_norm = None
                ppm = None

            setattr(self, f"reduced_embed{i + 1}", reduced_embed)
            setattr(self, f"pos_embed{i + 1}", pos_embed)
            setattr(self, f"block{i + 1}", block)
            setattr(self, f"block_attn{i + 1}", block_attn)
            setattr(self, f"ppm{i + 1}", ppm)
            setattr(self, f"ppm_layer_norm{i + 1}", ppm_layer_norm)

            if i == self.pyramid_stages - 1:
                if args.encoder_normalize_before:
                    self.layer_norm = LayerNorm(embed_dim)
                else:
                    self.layer_norm = None

        if self.use_ppm:
            self.ppm_weight = nn.Parameter(torch.Tensor(self.pyramid_stages).fill_(1.0))
            self.ppm_weight.data = self.ppm_weight.data / self.ppm_weight.data.sum(0, keepdim=True)

        self.use_ctc = "sate" in args.arch or \
                       (("ctc" in getattr(args, "criterion", False)) and
                        (getattr(args, "ctc_weight", False) > 0))
        if self.use_ctc:
            self.ctc_layer = (args.encoder_layers + args.ctc_layer) % args.encoder_layers
            self.inter_ctc = True if self.ctc_layer != args.encoder_layers else False

            if task.source_dictionary == task.target_dictionary and getattr(args, "share_all_embeddings", False):
                self.ctc_projection = nn.Linear(
                    embed_tokens.weight.shape[1],
                    embed_tokens.weight.shape[0],
                    bias=False,
                )
                self.ctc_projection.weight = embed_tokens.weight
            else:
                embed_dim = self.pyramid_embed_dims[-1]

                if self.inter_ctc:
                    ctc_layer = self.ctc_layer
                    for i in range(self.pyramid_stages):
                        ctc_layer -= self.pyramid_layers[i]
                        if ctc_layer <= 0:
                            embed_dim = self.pyramid_embed_dims[i]
                            break
                    self.ctc_layer_norm = LayerNorm(embed_dim)

                self.ctc_projection = nn.Linear(embed_dim, len(task.source_dictionary), bias=False)
                nn.init.normal_(
                    self.ctc_projection.weight, mean=0, std=embed_dim ** -0.5
                )
            self.ctc_dropout_module = FairseqDropout(
                p=args.dropout, module_name=self.__class__.__name__
            )
            self.softmax = nn.Softmax(dim=-1)

    def forward(self, src_tokens, src_lengths):

        batch = src_tokens.size(0)
        x = src_tokens.transpose(0, 1)
        input_lengths = src_lengths

        # padding to the multiply of 2
        max_len = x.size(0)
        length = reduce(lambda a, b: a*b, self.pyramid_sr_ratios)
        padding_to_len = (length - max_len % length)
        if padding_to_len > 0:
            padding_for_pyramid = x.new_zeros((padding_to_len, batch, x.size(2)))
            x = torch.cat([x, padding_for_pyramid], dim=0)

        layer_idx = 0
        ctc_logit = None
        prev_state = []
        prev_padding = []
        for i in range(self.pyramid_stages):
            reduced_embed = getattr(self, f"reduced_embed{i + 1}")
            pos_embed = getattr(self, f"pos_embed{i + 1}")
            block = getattr(self, f"block{i + 1}")
            block_attn = getattr(self, f"block_attn{i + 1}")

            # if i == 0:
            #    x = self.embed_scale * x

            x, input_lengths, encoder_padding_mask = reduced_embed(x, input_lengths)
            # add the position encoding and dropout
            if pos_embed:
                positions = pos_embed(encoder_padding_mask).transpose(0, 1)
                if self.attn_type != "rel_selfattn":
                    x += positions
                positions = self.dropout(positions)
            else:
                positions = None

            if i == 0:
                x = self.dropout(x)

            for layer in block:
                x = layer(x, encoder_padding_mask, pos_emb=positions)
                layer_idx += 1

                if self.use_ctc and self.inter_ctc and self.ctc_layer == layer_idx:
                    ctc_logit = self.ctc_layer_norm(x)

            prev_state.append(x)
            prev_padding.append(encoder_padding_mask)

            if block_attn is not None:
                x = block_attn(x, prev_state[-1], prev_padding[-1])

        if self.use_ppm:
            pool_state = [x]
            seq_len, bsz, dim = x.size()
            i = -1
            for state in prev_state[:-1]:
                i += 1
                ppm = getattr(self, f"ppm{i + 1}")
                ppm_layer_norm = getattr(self, f"ppm_layer_norm{i + 1}")

                state = ppm_layer_norm(state)
                state = state.permute(1, 2, 0)
                state = nn.functional.adaptive_avg_pool1d(state, seq_len)
                state = ppm(state)
                state = state.permute(2, 0, 1)
                pool_state.append(state)
            x = (torch.stack(pool_state, dim=0) * self.ppm_weight.view(-1, 1, 1, 1)).sum(0)

        if self.layer_norm is not None:
            x = self.layer_norm(x)

        return {
            "encoder_out": [x],  # T x B x C
            "encoder_padding_mask": [encoder_padding_mask],  # B x T
            "encoder_embedding": [],  # B x T x C
            "encoder_states": [],  # List[T x B x C]
            "ctc_logit": [ctc_logit if ctc_logit is not None else x],
            "src_tokens": [],
            "src_lengths": [],
        }

    def compute_ctc_logit(self, encoder_out):
        assert self.use_ctc, "CTC is not available!"

        if isinstance(encoder_out, dict) and "ctc_logit" in encoder_out:
            encoder_state = encoder_out["ctc_logit"][0]
        else:
            encoder_state = encoder_out
        ctc_logit = self.ctc_projection(self.ctc_dropout_module(encoder_state))

        return ctc_logit

    def compute_ctc_prob(self, encoder_out, temperature=1.0):
        assert self.use_ctc, "CTC is not available!"

        ctc_logit = self.compute_ctc_logit(encoder_out) / temperature

        return self.softmax(ctc_logit)

    def reorder_encoder_out(self, encoder_out, new_order):
        new_encoder_out = (
            [] if len(encoder_out["encoder_out"]) == 0
            else [x.index_select(1, new_order) for x in encoder_out["encoder_out"]]
        )

        new_encoder_padding_mask = (
            [] if len(encoder_out["encoder_padding_mask"]) == 0
            else [x.index_select(0, new_order) for x in encoder_out["encoder_padding_mask"]]
        )

        new_encoder_embedding = (
            [] if len(encoder_out["encoder_embedding"]) == 0
            else [x.index_select(0, new_order) for x in encoder_out["encoder_embedding"]]
        )

        encoder_states = encoder_out["encoder_states"]
        if len(encoder_states) > 0:
            for idx, state in enumerate(encoder_states):
                encoder_states[idx] = state.index_select(1, new_order)

        return {
            "encoder_out": new_encoder_out,  # T x B x C
            "encoder_padding_mask": new_encoder_padding_mask,  # B x T
            "encoder_embedding": new_encoder_embedding,  # B x T x C
            "encoder_states": encoder_states,  # List[T x B x C]
            "src_tokens": [],  # B x T
            "src_lengths": [],  # B x 1
        }


@register_model_architecture(model_name="pys2t_transformer", arch_name="pys2t_transformer")
def base_architecture(args):
    # Convolutional subsampler
    args.conv_kernel_sizes = getattr(args, "conv_kernel_sizes", "")
    args.conv_channels = getattr(args, "conv_channels", 1024)

    # Pyramid
    args.pyramid_stages = getattr(args, "pyramid_stages", None)
    args.pyramid_layers = getattr(args, "pyramid_layers", None)
    args.pyramid_sr_ratios = getattr(args, "pyramid_sr_ratios", None)
    args.pyramid_attn_sample_ratios = getattr(args, "pyramid_attn_sample_ratios", None)
    args.pyramid_embed_dims = getattr(args, "pyramid_embed_dims", None)
    args.pyramid_kernel_sizes = getattr(args, "pyramid_kernel_sizes", None)
    args.pyramid_ffn_ratios = getattr(args, "pyramid_ffn_ratios", None)
    args.pyramid_heads = getattr(args, "pyramid_heads", None)
    args.pyramid_position_embed = getattr(args, "pyramid_position_embed", None)
    args.pyramid_reduced_embed = getattr(args, "pyramid_reduced_embed", "conv")
    args.pyramid_embed_norm = getattr(args, "pyramid_embed_norm", False)

    args.ctc_layer = getattr(args, "ctc_layer", -1)

    args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 512)
    args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 2048)
    args.encoder_layers = getattr(args, "encoder_layers", 12)
    args.encoder_attention_type = getattr(args, "encoder_attention_type", "selfattn")
    args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 8)
    args.encoder_normalize_before = getattr(args, "encoder_normalize_before", True)
    args.decoder_embed_dim = getattr(args, "decoder_embed_dim", args.encoder_embed_dim)
    args.decoder_ffn_embed_dim = getattr(
        args, "decoder_ffn_embed_dim", args.encoder_ffn_embed_dim
    )
    args.decoder_layers = getattr(args, "decoder_layers", 6)
    args.decoder_attention_type = getattr(args, "decoder_attention_type", "selfattn")
    args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 8)
    args.decoder_normalize_before = getattr(args, "decoder_normalize_before", True)
    args.decoder_learned_pos = getattr(args, "decoder_learned_pos", False)
    args.dropout = getattr(args, "dropout", 0.1)
    args.attention_dropout = getattr(args, "attention_dropout", args.dropout)
    args.activation_dropout = getattr(args, "activation_dropout", args.dropout)
    args.activation_fn = getattr(args, "activation_fn", "relu")
    args.adaptive_softmax_cutoff = getattr(args, "adaptive_softmax_cutoff", None)
    args.adaptive_softmax_dropout = getattr(args, "adaptive_softmax_dropout", 0)
    args.share_decoder_input_output_embed = getattr(
        args, "share_decoder_input_output_embed", False
    )
    args.no_token_positional_embeddings = getattr(
        args, "no_token_positional_embeddings", False
    )
    args.adaptive_input = getattr(args, "adaptive_input", False)
    args.decoder_layerdrop = getattr(args, "decoder_layerdrop", 0.0)
    args.decoder_output_dim = getattr(
        args, "decoder_output_dim", args.decoder_embed_dim
    )
    args.decoder_input_dim = getattr(args, "decoder_input_dim", args.decoder_embed_dim)
    args.no_scale_embedding = getattr(args, "no_scale_embedding", False)
    args.quant_noise_pq = getattr(args, "quant_noise_pq", 0)

    args.max_encoder_relative_length = getattr(args, 'max_encoder_relative_length', -1)
    args.max_decoder_relative_length = getattr(args, 'max_decoder_relative_length', -1)
    args.k_only = getattr(args, 'k_only', True)


@register_model_architecture("pys2t_transformer", "pys2t_transformer_s")
def pys2t_transformer_s(args):
    args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 256)
    args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 256 * 8)
    args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 4)
    args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 4)
    args.dropout = getattr(args, "dropout", 0.1)

    args.pyramid_stages = getattr(args, "pyramid_stages", 4)
    args.pyramid_layers = getattr(args, "pyramid_layers", "3_3_3_3")
    args.pyramid_embed_dims = getattr(args, "pyramid_embed_dims", "64_128_256_512")
    args.pyramid_kernel_sizes = getattr(args, "pyramid_kernel_sizes", "2_2_2_2")
    args.pyramid_ffn_ratios = getattr(args, "pyramid_ffn_ratios", "4_4_4_4")
    args.pyramid_attn_sample_ratios = getattr(args, "pyramid_attn_sample_ratios", "8_4_2_1")
    args.pyramid_sr_ratios = getattr(args, "pyramid_sr_ratios", "2_2_2_2")
    args.pyramid_heads = getattr(args, "pyramid_heads", "1_2_4_8")
    args.pyramid_position_embed = getattr(args, "pyramid_position_embed", "1_1_1_1")
    args.pyramid_reduced_embed = getattr(args, "pyramid_reduced_embed", "conv")
    args.pyramid_embed_norm = getattr(args, "pyramid_embed_norm", False)

    base_architecture(args)


@register_model_architecture("pys2t_transformer", "pys2t_transformer_s_relative")
def pys2t_transformer_s_relative(args):
    args.max_encoder_relative_length = 100
    args.max_decoder_relative_length = 20
    args.k_only = True
    pys2t_transformer_s(args)


@register_model_architecture("pys2t_transformer", "pys2t_transformer_xs")
def pys2t_transformer_xs(args):
    args.encoder_layers = getattr(args, "encoder_layers", 6)
    args.decoder_layers = getattr(args, "decoder_layers", 3)
    args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 256 * 4)
    args.dropout = getattr(args, "dropout", 0.3)
    pys2t_transformer_s(args)


@register_model_architecture("pys2t_transformer", "pys2t_transformer_sp")
def pys2t_transformer_sp(args):
    args.encoder_layers = getattr(args, "encoder_layers", 16)
    pys2t_transformer_s(args)


@register_model_architecture("pys2t_transformer", "pys2t_transformer_m")
def pys2t_transformer_m(args):
    args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 512)
    args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 512 * 4)
    args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 8)
    args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 8)
    args.dropout = getattr(args, "dropout", 0.15)
    base_architecture(args)


@register_model_architecture("pys2t_transformer", "pys2t_transformer_mp")
def pys2t_transformer_mp(args):
    args.encoder_layers = getattr(args, "encoder_layers", 16)
    pys2t_transformer_m(args)


@register_model_architecture("pys2t_transformer", "pys2t_transformer_l")
def pys2t_transformer_l(args):
    args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 1024)
    args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 1024 * 4)
    args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 16)
    args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 16)
    args.dropout = getattr(args, "dropout", 0.2)
    base_architecture(args)


@register_model_architecture("pys2t_transformer", "pys2t_transformer_lp")
def pys2t_transformer_lp(args):
    args.encoder_layers = getattr(args, "encoder_layers", 16)
    pys2t_transformer_l(args)