fairseq/models/speech_to_text/xm_transformer_unity.py

# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import copy
import logging

from fairseq.models import (
    FairseqEncoderModel,
    FairseqLanguageModel,
    register_model,
    register_model_architecture,
)
from fairseq.models.speech_to_speech.modules.ctc_decoder import CTCDecoder
from fairseq.models.speech_to_speech.modules.transformer_encoder import (
    TransformerEncoderNoEmb,
)
from fairseq.models.speech_to_text.xm_transformer import XMTransformerModel
from fairseq.models.speech_to_text.xm_transformer import (
    base_architecture as xm_t_base_architecture,
)
from fairseq.models.speech_to_text.xm_transformer import (
    build_embedding,
    need_finetuning,
    set_default_adaptor_args,
    set_default_general_args,
    set_default_transformer_decoder_args,
    set_default_w2v_encoder_args,
)
from fairseq.models.transformer import Linear, TransformerDecoder, TransformerModelBase
from fairseq.models.transformer.transformer_decoder_aug import AugTransformerDecoder

logger = logging.getLogger(__name__)


def unit_transformer_decoder_arch_base(
    args, decoder_layers=6, decoder_embed_dim=768, decoder_attention_heads=12
):
    args.encoder_layers = decoder_layers
    args.decoder_layers = decoder_layers
    args.decoder_embed_dim = decoder_embed_dim
    args.decoder_ffn_embed_dim = decoder_embed_dim * 4
    args.decoder_attention_heads = decoder_attention_heads
    args.encoder_embed_dim = args.decoder_embed_dim
    args.decoder_output_dim = decoder_embed_dim
    args.decoder_input_dim = decoder_embed_dim


def unit_transformer_decoder_arch_large(
    args, decoder_layers=12, decoder_embed_dim=1024, decoder_attention_heads=16
):
    args.encoder_layers = decoder_layers
    args.decoder_layers = decoder_layers
    args.decoder_embed_dim = decoder_embed_dim
    args.decoder_ffn_embed_dim = decoder_embed_dim * 4
    args.decoder_attention_heads = decoder_attention_heads
    args.encoder_embed_dim = args.decoder_embed_dim
    args.decoder_output_dim = decoder_embed_dim
    args.decoder_input_dim = decoder_embed_dim


@register_model("unity_xm_transformer")
class XMTransformerModelUnitY(XMTransformerModel):
    @classmethod
    def hub_models(cls):
        base_url = "http://dl.fbaipublicfiles.com/fairseq/s2t"
        model_ids = []
        return {i: f"{base_url}/{i}.tar.gz" for i in model_ids}

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

    @classmethod
    def add_args(cls, parser):
        """Add model-specific arguments to the parser."""
        XMTransformerModel.add_args(parser)
        parser.add_argument(
            "--translation-decoder-layers",
            type=int,
            default=4,
            metavar="N",
            help="num decoder layers in the first-pass translation module",
        )
        parser.add_argument(
            "--synthesizer-encoder-layers",
            type=int,
            default=0,
            metavar="N",
            help="num encoder layers in the second-pass synthesizer module",
        )
        parser.add_argument(
            "--synthesizer-augmented-cross-attention",
            action="store_true",
            default=False,
            help="augmented cross-attention over speech encoder output",
        )
        parser.add_argument(
            "--load-pretrained-aux-decoder-from",
            type=str,
            metavar="STR",
            help="model to take decoder weights from (for initialization)",
        )

    @classmethod
    def build_text_decoder(cls, args, tgt_dict):
        _args = copy.deepcopy(args)

        if args.adaptor_proj or args.encoder_proj:  # not V0 arch
            _args.encoder_embed_dim = _args.decoder_embed_dim
        _args.dropout = args.decoder_dropout
        _args.attention_dropout = args.decoder_attention_dropout
        _args.activation_dropout = args.decoder_activation_dropout
        _args.layerdrop = _args.decoder_layerdrop
        _args.decoder_layers = _args.translation_decoder_layers

        embed_tokens = build_embedding(tgt_dict, _args.decoder_embed_dim)
        decoder = TransformerDecoder(_args, tgt_dict, embed_tokens)

        if getattr(args, "load_pretrained_aux_decoder_from", None) is not None:
            decoder = cls.maybe_load_pretrained(
                decoder, getattr(args, "load_pretrained_aux_decoder_from", None)
            )

            for k, p in decoder.named_parameters():
                p.requires_grad = need_finetuning(args.finetune_decoder_params, k)
        return decoder

    @classmethod
    def build_decoder(cls, args, task, aug_attn=False):
        _args = copy.deepcopy(args)
        _args.layerdrop = 0.0  # turn off layerdrop for shallow layers

        _args.encoder_embed_dim = args.decoder_embed_dim

        proj = None
        if args.decoder_embed_dim != _args.decoder_embed_dim:
            proj = Linear(args.decoder_embed_dim, _args.decoder_embed_dim)

        embed_tokens = build_embedding(task.target_dictionary, _args.decoder_embed_dim)
        decoder_cls = AugTransformerDecoder if aug_attn else TransformerDecoder
        decoder = decoder_cls(_args, task.target_dictionary, embed_tokens)

        if getattr(args, "load_pretrained_decoder_from", None) is not None:
            # load all layers first and then discard the bottom layers
            embed_tokens = build_embedding(
                task.target_dictionary, _args.decoder_embed_dim
            )
            decoder_tmp = decoder_cls(_args, task.target_dictionary, embed_tokens)
            decoder_tmp = cls.maybe_load_pretrained(
                decoder_tmp, getattr(_args, "load_pretrained_decoder_from", None)
            )
            state_dict = decoder_tmp.state_dict()
            for k, p in decoder.named_parameters():
                p.data = state_dict[k].data
                p.requires_grad = need_finetuning(_args.finetune_decoder_params, k)
            decoder.layers = decoder.layers[-_args.decoder_layers :]

        return decoder, proj, _args

    @classmethod
    def build_model(cls, args, task):
        """Build a new model instance."""

        # make sure all arguments are present in older models
        xm_t_base_architecture(args)

        encoder = cls.build_encoder(args)
        decoder, proj, unit_args = cls.build_decoder(
            args,
            task,
            aug_attn=getattr(args, "synthesizer_augmented_cross_attention", False),
        )
        base_model = cls(encoder, decoder)
        setattr(base_model, "proj", proj)

        base_model.t2u_augmented_cross_attn = getattr(
            args, "synthesizer_augmented_cross_attention", False
        )

        # set up multitask decoders
        base_model.mt_task_name = None
        base_model.multitask_decoders = {}
        has_first_pass_decoder = False
        for task_name, task_obj in task.multitask_tasks.items():
            if task_obj.is_first_pass_decoder:
                has_first_pass_decoder = True
                base_model.mt_task_name = task_name

            task_decoder = cls.build_multitask_decoder(
                args,
                task_obj.args,
                task_obj.target_dictionary,
                args.decoder_embed_dim,
                task_obj.is_first_pass_decoder,
            )

            setattr(base_model, f"{task_name}_decoder", task_decoder)
            decoder_model_cls = (
                FairseqEncoderModel
                if task_obj.args.decoder_type == "ctc"
                else FairseqLanguageModel
            )
            base_model.multitask_decoders[task_name] = decoder_model_cls(
                getattr(base_model, f"{task_name}_decoder")
            )

        assert has_first_pass_decoder, "set at least one intermediate non-CTC decoder"

        # set up encoder on top of the auxiliary MT decoder
        if getattr(args, "synthesizer_encoder_layers", 0) > 0:
            base_model.synthesizer_encoder = cls.build_t2u_encoder(unit_args)
        else:
            base_model.synthesizer_encoder = None

        return base_model

    @classmethod
    def build_t2u_encoder(cls, args):
        _args = copy.deepcopy(args)
        _args.encoder_layers = _args.synthesizer_encoder_layers
        _args.encoder_embed_dim = args.decoder_embed_dim
        _args.encoder_ffn_embed_dim = args.decoder_ffn_embed_dim
        _args.encoder_attention_heads = args.decoder_attention_heads
        _args.encoder_normalize_before = True
        return TransformerEncoderNoEmb(_args)

    def forward(
        self,
        src_tokens,
        src_lengths,
        prev_output_tokens,
        prev_output_tokens_mt,
        return_all_hiddens=False,
        tgt_speaker=None,
        **kwargs,
    ):
        """
        The forward method inherited from the base class has a **kwargs
        argument in its input, which is not supported in torchscript. This
        method overwrites the forward method definition without **kwargs.
        """
        encoder_out = self.encoder(
            src_tokens=src_tokens, src_lengths=src_lengths, **kwargs
        )

        # 1. MT decoder
        mt_decoder = getattr(self, f"{self.mt_task_name}_decoder")
        mt_decoder_out = mt_decoder(
            prev_output_tokens_mt,
            encoder_out=encoder_out,
        )
        x = mt_decoder_out[1]["inner_states"][-1]
        if mt_decoder.layer_norm is not None:
            x = mt_decoder.layer_norm(x)
        if self.proj is not None:
            x = self.proj(x)

        mt_decoder_padding_mask = None
        if prev_output_tokens_mt.eq(mt_decoder.padding_idx).any():
            mt_decoder_padding_mask = prev_output_tokens_mt.eq(mt_decoder.padding_idx)

        # 2. T2U encoder
        if self.synthesizer_encoder is not None:
            t2u_encoder_out = self.synthesizer_encoder(
                x,
                mt_decoder_padding_mask,
            )
        else:
            t2u_encoder_out = {
                "encoder_out": [x],  # T x B x C
                "encoder_padding_mask": [mt_decoder_padding_mask],  # B x T
            }

        # 3. T2U decoder
        if self.t2u_augmented_cross_attn:
            decoder_out = self.decoder(
                prev_output_tokens,
                encoder_out=encoder_out,
                encoder_out_aug=t2u_encoder_out,
            )
        else:
            decoder_out = self.decoder(
                prev_output_tokens,
                encoder_out=t2u_encoder_out,
            )
        if return_all_hiddens:
            decoder_out[-1]["encoder_states"] = encoder_out["encoder_out"]
            # NOTE: from the top layer
            decoder_out[-1]["encoder_padding_mask"] = encoder_out[
                "encoder_padding_mask"
            ]
        decoder_out[-1]["mt_decoder_out"] = mt_decoder_out
        return decoder_out


@register_model_architecture(
    model_name="unity_xm_transformer", arch_name="unity_xm_transformer"
)
def base_architecture_unity(args):
    set_default_general_args(args)
    set_default_w2v_encoder_args(args)
    set_default_adaptor_args(args)
    set_default_transformer_decoder_args(args)

    args.layernorm_embedding = False
    args.decoder_learned_pos = False


# for old models
@register_model_architecture(
    model_name="unity_xm_transformer", arch_name="xm_transformer_t2"
)
def base_architecture_unity_legacy(args):
    base_architecture_unity(args)