mmpretrain/models/multimodal/ofa/ofa_modules.py

# Copyright (c) OpenMMLab. All rights reserved.
import math
from dataclasses import dataclass
from functools import partial
from typing import List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from mmcv.cnn.bricks import DropPath
from mmengine.model import BaseModule
from mmengine.utils import digit_version
from transformers.modeling_outputs import (
    BaseModelOutputWithPastAndCrossAttentions, ModelOutput, Seq2SeqLMOutput)
from transformers.modeling_utils import (GenerationConfig, GenerationMixin,
                                         PretrainedConfig)

from mmpretrain.registry import MODELS
from ...backbones.resnet import Bottleneck, ResNet

if digit_version(torch.__version__) >= digit_version('1.10.0'):
    torch_meshgrid = partial(torch.meshgrid, indexing='ij')
else:
    torch_meshgrid = torch.meshgrid


def make_token_bucket_position(bucket_size, max_position=1024):
    context_pos = torch.arange(max_position, dtype=torch.long)[:, None]
    memory_pos = torch.arange(max_position, dtype=torch.long)[None, :]
    relative_pos = context_pos - memory_pos
    sign = torch.sign(relative_pos)
    mid = bucket_size // 2
    abs_pos = torch.where((relative_pos < mid) & (relative_pos > -mid),
                          mid - 1, torch.abs(relative_pos))
    log_pos = torch.ceil(
        torch.log(abs_pos / mid) / math.log(
            (max_position - 1) / mid) * (mid - 1)) + mid
    log_pos = log_pos.int()
    bucket_pos = torch.where(abs_pos.le(mid), relative_pos,
                             log_pos * sign).long()
    return bucket_pos + bucket_size - 1


def make_image_bucket_position(bucket_size, num_relative_distance):
    coords_h = torch.arange(bucket_size)
    coords_w = torch.arange(bucket_size)
    # (2, h, w)
    coords = torch.stack(torch_meshgrid([coords_h, coords_w]))
    # (2, h*w)
    coords_flatten = torch.flatten(coords, 1)
    # (2, h*w, h*w)
    relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
    # (h*w, h*w, 2)
    relative_coords = relative_coords.permute(1, 2, 0).contiguous()
    relative_coords[:, :, 0] += bucket_size - 1  # shift to start from 0
    relative_coords[:, :, 1] += bucket_size - 1
    relative_coords[:, :, 0] *= 2 * bucket_size - 1
    relative_position_index = torch.zeros(
        size=(bucket_size * bucket_size + 1, ) * 2,
        dtype=relative_coords.dtype)
    # (h*w, h*w)
    relative_position_index[1:, 1:] = relative_coords.sum(-1)
    relative_position_index[0, 0:] = num_relative_distance - 3
    relative_position_index[0:, 0] = num_relative_distance - 2
    relative_position_index[0, 0] = num_relative_distance - 1
    return relative_position_index


def _make_causal_mask(input_ids_shape: torch.Size,
                      dtype: torch.dtype,
                      past_key_values_length: int = 0):
    """Make causal mask used for uni-directional self-attention."""
    bsz, tgt_len = input_ids_shape
    mask = torch.full((tgt_len, tgt_len), float('-inf'))
    mask_cond = torch.arange(mask.size(-1))
    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
    mask = mask.to(dtype)

    if past_key_values_length > 0:
        mask = torch.cat(
            [torch.zeros(tgt_len, past_key_values_length, dtype=dtype), mask],
            dim=-1)
    return mask[None, None, :, :].expand(bsz, 1, tgt_len,
                                         tgt_len + past_key_values_length)


def _expand_mask(mask: torch.Tensor,
                 dtype: torch.dtype,
                 tgt_len: Optional[int] = None):
    """Expands attention_mask from ``[B, L_s]`` to ``[B, 1, L_t, L_s]``.

    Where ``B`` is batch_size, `L_s`` is the source sequence length, and
    ``L_t`` is the target sequence length.
    """
    bsz, src_len = mask.size()
    tgt_len = tgt_len if tgt_len is not None else src_len

    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len,
                                                  src_len).to(dtype)
    return expanded_mask.masked_fill(expanded_mask.bool(),
                                     torch.finfo(dtype).min)


class MultiheadAttention(BaseModule):
    """Multi-head Attention Module for OFA.

    Args:
        embedding_dim (int): The embedding dimension of query.
        num_heads (int): Parallel attention heads.
        kdim (int, optional): The embedding dimension of key.
            Defaults to None, which means the same as the `embedding_dim`.
        vdim (int, optional): The embedding dimension of value.
            Defaults to None, which means the same as the `embedding_dim`.
        attn_drop (float): Dropout rate of the dropout layer after the
            attention calculation of query and key. Defaults to 0.
        qkv_bias (bool): If True, add a learnable bias to q, k, v.
            Defaults to True.
        scale_factor (float): The scale of qk will be
            ``(head_dim * scale_factor) ** -0.5``. Defaults to 1.
        proj_bias (bool) If True, add a learnable bias to output projection.
            Defaults to True.
        init_cfg (dict, optional): The Config for initialization.
            Defaults to None.
    """

    def __init__(self,
                 embedding_dim,
                 num_heads,
                 kdim=None,
                 vdim=None,
                 attn_drop=0.,
                 scale_factor=1.,
                 qkv_bias=True,
                 proj_bias=True,
                 scale_heads=False,
                 init_cfg=None):
        super(MultiheadAttention, self).__init__(init_cfg=init_cfg)

        self.embedding_dim = embedding_dim
        self.num_heads = num_heads
        self.kdim = kdim or embedding_dim
        self.vdim = vdim or embedding_dim

        self.head_dim = embedding_dim // num_heads
        self.scale = (self.head_dim * scale_factor)**-0.5

        self.q_proj = nn.Linear(embedding_dim, embedding_dim, bias=qkv_bias)
        self.k_proj = nn.Linear(self.kdim, embedding_dim, bias=qkv_bias)
        self.v_proj = nn.Linear(self.vdim, embedding_dim, bias=qkv_bias)
        self.out_proj = nn.Linear(embedding_dim, embedding_dim, bias=proj_bias)

        self.attn_drop = nn.Dropout(p=attn_drop)

        if scale_heads:
            self.c_attn = nn.Parameter(torch.ones(num_heads))
        else:
            self.c_attn = None

    def forward(
        self,
        query,
        key_value=None,
        attn_mask=None,
        attn_bias=None,
        past_key_value=None,
        output_attentions=False,
    ):
        B, _, C = query.shape
        assert C == self.head_dim * self.num_heads

        is_cross_attention = key_value is not None
        if key_value is None:
            key_value = query

        # (B, L, C) -> (B, num_heads, L, head_dims)
        q = self.q_proj(query).reshape(B, -1, self.num_heads,
                                       self.head_dim).transpose(1, 2)

        if is_cross_attention and past_key_value is not None:
            # Reuse key and value in cross_attentions
            k, v = past_key_value
        else:
            k = self.k_proj(key_value).reshape(B, -1, self.num_heads,
                                               self.head_dim).transpose(1, 2)
            v = self.v_proj(key_value).reshape(B, -1, self.num_heads,
                                               self.head_dim).transpose(1, 2)
            if past_key_value is not None:
                past_key, past_value = past_key_value
                k = torch.cat([past_key, k], dim=2)
                v = torch.cat([past_value, v], dim=2)

        past_key_value = (k, v)

        attn_weights = q @ k.transpose(-2, -1) * self.scale

        if attn_bias is not None:
            src_len = k.size(2)
            attn_weights[:, :, -src_len:] += attn_bias[:, :, -src_len:]

        if attn_mask is not None:
            attn_weights += attn_mask
        attn_weights = torch.softmax(attn_weights, dim=-1)
        attn = self.attn_drop(attn_weights) @ v

        if self.c_attn is not None:
            attn = torch.einsum('bhlc,h->bhlc', attn, self.c_attn)

        # (B, num_heads, L, head_dims) -> (B, L, C)
        attn = attn.transpose(1, 2).reshape(B, -1, self.embedding_dim)
        attn = self.out_proj(attn)

        if output_attentions:
            return attn, attn_weights, past_key_value
        else:
            return attn, None, past_key_value


@MODELS.register_module(force=True)
class OFAResNet(ResNet):
    """ResNet module for OFA.

    The ResNet in OFA has only three stages.
    """
    arch_settings = {
        50: (Bottleneck, (3, 4, 6)),
        101: (Bottleneck, (3, 4, 23)),
        152: (Bottleneck, (3, 8, 36)),
    }

    def __init__(self, depth, *args, **kwargs):
        super().__init__(
            depth=depth,
            *args,
            num_stages=3,
            out_indices=(2, ),
            dilations=(1, 1, 1),
            strides=(1, 2, 2),
            **kwargs)


@dataclass
class OFAEncoderOutput(ModelOutput):
    """OFA encoder outputs.

    Args:
        last_hidden_state (torch.tensor): The hidden-states of the output at
            the last layer of the model. The shape is (B, L, C).
        hidden_states (Tuple[torch.tensor]): The initial embedding and the
            output of each layer. The shape of every item is (B, L, C).
        attentions (Tuple[torch.tensor]): The attention weights after the
            attention softmax, used to compute the weighted average in the
            self-attention heads. The shape of every item is
            (B, num_heads, L, L).
        position_embedding (torch.tensor): The positional embeddings of the
            inputs. The shape is (B, L, C).
    """

    last_hidden_state: torch.FloatTensor = None
    padding_mask: torch.Tensor = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None
    position_embedding: Optional[torch.FloatTensor] = None


class OFAEncoderLayer(nn.Module):
    """OFAEncoder layer block."""

    def __init__(self,
                 embedding_dim,
                 num_heads,
                 dropout_rate=0.,
                 drop_path_rate=0.,
                 attn_drop=0.,
                 act_drop=0.,
                 scale_factor=2.,
                 mlp_ratio=4.,
                 scale_heads=True,
                 normformer=True,
                 pre_norm=True,
                 act_cfg=dict(type='GELU')):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.pre_norm = pre_norm

        self.attn = MultiheadAttention(
            embedding_dim=embedding_dim,
            num_heads=num_heads,
            attn_drop=attn_drop,
            scale_factor=scale_factor,
            scale_heads=scale_heads,
        )

        mid_channels = int(embedding_dim * mlp_ratio)
        self.fc1 = nn.Linear(embedding_dim, mid_channels)
        self.fc2 = nn.Linear(mid_channels, embedding_dim)
        self.act = MODELS.build(act_cfg)
        self.act_drop = nn.Dropout(
            act_drop) if act_drop > 0. else nn.Identity()

        # LayerNorm between attention block and ffn block.
        self.attn_ln = nn.LayerNorm(embedding_dim)
        self.ffn_ln = nn.LayerNorm(embedding_dim)

        # Extra LayerNorm
        self.normformer = normformer
        if self.normformer:
            self.attn_mid_ln = nn.LayerNorm(embedding_dim)
            self.ffn_mid_ln = nn.LayerNorm(mid_channels)

        self.dropout = nn.Dropout(dropout_rate)
        self.drop_path = DropPath(
            drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()

    def forward(self,
                x,
                attention_mask=None,
                attn_bias=None,
                output_attentions=False):
        """Forward the encoder layer.

        Args:
            x (torch.tensor): The input to the layer of shape ``(B, L, C)``.
            attention_mask (torch.Tensor, optional): The attention mask of size
                ``(B, 1, L, L)``, where padding elements are indicated by very
                large negative values. Defaults to None.
            attn_bias (torch.tensor, optional): The bias for positional
                information. Defaults to None.
            output_attentions (bool): Whether to return the attentions tensors
                of the attention layer.

        Returns:
            List[torch.tensor]: The first element is the encoded output of
            shape ``(B, L, C)``. And the second element is the output
            attentions if ``output_attentions=True``.
        """
        residual = x

        # Attention block
        if self.pre_norm:
            x = self.attn_ln(x)
        x, attn_weights, _ = self.attn(
            query=x,
            attn_mask=attention_mask,
            attn_bias=attn_bias,
            output_attentions=output_attentions)
        if self.normformer:
            x = self.attn_mid_ln(x)
        x = self.dropout(x)
        x = residual + self.drop_path(x)
        if not self.pre_norm:
            x = self.attn_ln(x)

        residual = x

        # FFN block
        if self.pre_norm:
            x = self.ffn_ln(x)
        x = self.act(self.fc1(x))
        x = self.act_drop(x)
        if self.normformer:
            x = self.ffn_mid_ln(x)
        x = self.fc2(x)
        x = self.dropout(x)
        x = residual + self.drop_path(x)
        if not self.pre_norm:
            x = self.ffn_ln(x)

        if output_attentions:
            return [x, attn_weights]
        else:
            return [x]


class OFADecoderLayer(nn.Module):
    """OFADecoder layer block."""

    def __init__(self,
                 embedding_dim,
                 num_heads,
                 dropout_rate=0.,
                 drop_path_rate=0.,
                 attn_drop=0.,
                 act_drop=0.,
                 scale_factor=2.,
                 mlp_ratio=4.,
                 encoder_embed_dim=None,
                 scale_heads=True,
                 normformer=True,
                 pre_norm=True,
                 act_cfg=dict(type='GELU')):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.pre_norm = pre_norm

        self.self_attn = MultiheadAttention(
            embedding_dim=embedding_dim,
            num_heads=num_heads,
            attn_drop=attn_drop,
            scale_factor=scale_factor,
            scale_heads=scale_heads,
        )

        self.cross_attn = MultiheadAttention(
            embedding_dim=embedding_dim,
            kdim=encoder_embed_dim,
            vdim=encoder_embed_dim,
            num_heads=num_heads,
            attn_drop=attn_drop,
            scale_factor=scale_factor,
            scale_heads=scale_heads,
        )

        mid_channels = int(embedding_dim * mlp_ratio)
        self.fc1 = nn.Linear(embedding_dim, mid_channels)
        self.fc2 = nn.Linear(mid_channels, embedding_dim)
        self.act = MODELS.build(act_cfg)
        self.act_drop = nn.Dropout(
            act_drop) if act_drop > 0. else nn.Identity()

        # LayerNorm between attention block and ffn block.
        self.self_attn_ln = nn.LayerNorm(embedding_dim)
        self.cross_attn_ln = nn.LayerNorm(embedding_dim)
        self.ffn_ln = nn.LayerNorm(embedding_dim)

        # Extra LayerNorm
        self.normformer = normformer
        if self.normformer:
            self.self_attn_mid_ln = nn.LayerNorm(embedding_dim)
            self.cross_attn_mid_ln = nn.LayerNorm(embedding_dim)
            self.ffn_mid_ln = nn.LayerNorm(mid_channels)

        self.dropout = nn.Dropout(dropout_rate)
        self.drop_path = DropPath(
            drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()

    def forward(
        self,
        x,
        attention_mask=None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.Tensor] = None,
        past_key_value: Optional[List[torch.Tensor]] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        self_attn_bias: Optional[torch.Tensor] = None,
        cross_attn_bias: Optional[torch.Tensor] = None,
    ):
        """Forward the decoder layer.

        Args:
            x (torch.tensor): The input to the layer of shape ``(B, L, C)``.
            attention_mask (torch.Tensor, optional): The attention mask of size
                ``(B, 1, L, L)``, where padding elements are indicated by very
                large negative values. Defaults to None.
            encoder_hidden_states (torch.Tensor, optional): The cross attention
                input to the layer of size ``(B, L, C)``. Defaults to None.
            encoder_attention_mask (torch.Tensor, optional): The cross
                attention mask where padding elements are indicated by very
                large negative values. Defaults to None.
            past_key_value (Tuple[torch.tensor], optional): The cached past key
                and value projection states. Defaults to none.
            output_attentions (bool): whether to return the attentions tensors
                of all attention layers. Defaults to False.
            use_cache (bool, optional): Whether to use cache.
                Defaults to False.
            self_attn_bias (torch.Tensor, optional): The self attention bias
                for positional information. Defaults to None.
            cross_attn_bias (torch.Tensor, optional): The cross attention bias
                for positional information. Defaults to None.

        Returns:
            List[torch.tensor]: The first element is the encoded output of
            shape ``(B, L, C)``. The following two elements can be the output
            self-attentions and cross-attentions if ``output_attentions=True``.
            The following one element can be the cached past key and value
            projection states.
        """
        residual = x

        if past_key_value is not None:
            self_past_key_value = past_key_value[:2]
            cross_past_key_value = past_key_value[2:]
        else:
            self_past_key_value, cross_past_key_value = None, None

        # Self-Attention block
        if self.pre_norm:
            x = self.self_attn_ln(x)
        x, self_attn_weights, present_key_value = self.self_attn(
            query=x,
            past_key_value=self_past_key_value,
            attn_mask=attention_mask,
            output_attentions=output_attentions,
            attn_bias=self_attn_bias,
        )
        if self.normformer:
            x = self.self_attn_mid_ln(x)
        x = self.dropout(x)
        x = residual + self.drop_path(x)
        if not self.pre_norm:
            x = self.self_attn_ln(x)

        # Cross-Attention block
        if encoder_hidden_states is not None:
            residual = x
            if self.pre_norm:
                x = self.cross_attn_ln(x)
            x, cross_attn_weights, cross_key_value = self.cross_attn.forward(
                query=x,
                key_value=encoder_hidden_states,
                attn_mask=encoder_attention_mask,
                past_key_value=cross_past_key_value,
                output_attentions=output_attentions,
                attn_bias=cross_attn_bias)
            if self.normformer:
                x = self.cross_attn_mid_ln(x)
            x = self.dropout(x)
            x = residual + self.drop_path(x)
            if not self.pre_norm:
                x = self.cross_attn_ln(x)

            present_key_value = present_key_value + cross_key_value

        residual = x

        # FFN block
        if self.pre_norm:
            x = self.ffn_ln(x)
        x = self.act(self.fc1(x))
        x = self.act_drop(x)
        if self.normformer:
            x = self.ffn_mid_ln(x)
        x = self.fc2(x)
        x = self.dropout(x)
        x = residual + self.drop_path(x)
        if not self.pre_norm:
            x = self.ffn_ln(x)

        outputs = [x]

        if output_attentions:
            outputs.extend([self_attn_weights, cross_attn_weights])

        if use_cache:
            outputs.append(present_key_value)

        return outputs


class OFAEncoder(BaseModule):
    """The encoder module of OFA.

    Args:
        embed_tokens (nn.Embedding): The embedding module to embed the
            input tokens.
        embed_images (dict | nn.Module): The module to embed the input
            images into features. The output number of channels should
            be 1024.
        num_layers (int): The number of encoder layers. Defaults to 6.
        num_heads (int): The number of heads of attention. Defaults to 12.
        dropout_rate (float): The prob of dropout for embedding and
            transformer layers. Defaults to 0.
        drop_path_rate (float): The prob of droppath for transformer layers.
            Defaults to 0.
        max_source_positions (int): The maximum length of the input tokens.
            Defaults to 1024.
        token_bucket_size (int): The token bucket size, it's used as the
            maximum relative position index in relative position embedding
            of input tokens. Defaults to 256.
        image_bucket_size (int): The image bucket size, it's used to generate
            the image relative position embedding table. It should be larger
            than the shape of image feature map. Defaults to 42.
        attn_scale_factor (float): The scale factor to calculate qk scale in
            attentions. Defaults to 2.
        scale_embedding (bool): Whether to scale the embeddings by the square
            root of the dimension. Defaults to False.
        add_embedding_ln (bool): Whether to add an extra layer norm for token
            embeddings. Defaults to True.
        add_image_embedding_ln (bool): Whether to add an extra layer norm for
            image embeddings. Defaults to True.
        pre_norm (bool): Whether to do layer norm before attention and ffn
            blocks in transformer layers. Defaults to True.
        entangle_position_embedding (bool): Whether to add the position
            embedding on the embeddings directly. Defaults to False.
        init_cfg (dict, optional): The initialization config. Defaults to None.
    """

    def __init__(
        self,
        embed_tokens,
        embed_images: dict,
        num_layers=6,
        num_heads=12,
        dropout_rate=0.,
        drop_path_rate=0.,
        max_source_positions=1024,
        token_bucket_size=256,
        image_bucket_size=42,
        attn_scale_factor=2.,
        scale_embedding=False,
        add_embedding_ln=True,
        add_type_embed=True,
        add_image_embedding_ln=True,
        pre_norm=True,
        entangle_position_embedding=False,
        init_cfg=None,
    ):
        super().__init__(init_cfg=init_cfg)

        self.num_layers = num_layers
        embedding_dim = embed_tokens.embedding_dim
        self.embedding_dim = embedding_dim
        self.padding_idx = embed_tokens.padding_idx
        self.max_source_positions = max_source_positions
        self.num_heads = num_heads

        # Build embedding process components
        self.embed_tokens = embed_tokens
        self.embedding_scale = math.sqrt(
            embedding_dim) if scale_embedding else 1.0

        if not isinstance(embed_images, nn.Module):
            self.embed_images = MODELS.build(embed_images)
        else:
            self.embed_images = embed_images
        self.image_proj = nn.Linear(1024, embedding_dim)

        if add_embedding_ln:
            self.embedding_ln = nn.LayerNorm(embedding_dim)
        else:
            self.embedding_ln = None

        if add_type_embed:
            self.embed_type = nn.Embedding(2, embedding_dim)
        else:
            self.embed_type = None

        if add_image_embedding_ln:
            self.image_embedding_ln = nn.LayerNorm(embedding_dim)
        else:
            self.image_embedding_ln = None

        self.entangle_position_embedding = entangle_position_embedding

        # Build position embedding
        self.embed_positions = nn.Embedding(self.max_source_positions + 2,
                                            embedding_dim)
        self.pos_ln = nn.LayerNorm(embedding_dim)
        self.embed_image_positions = nn.Embedding(image_bucket_size**2 + 1,
                                                  embedding_dim)
        self.image_pos_ln = nn.LayerNorm(embedding_dim)

        self.pos_scaling = float(embedding_dim / num_heads *
                                 attn_scale_factor)**-0.5
        self.pos_q_linear = nn.Linear(embedding_dim, embedding_dim)
        self.pos_k_linear = nn.Linear(embedding_dim, embedding_dim)

        self.dropout = nn.Dropout(
            dropout_rate) if dropout_rate > 0. else nn.Identity()

        # Register token relative position embedding table
        self.token_bucket_size = token_bucket_size
        token_num_rel_dis = 2 * token_bucket_size - 1
        token_rp_bucket = make_token_bucket_position(token_bucket_size,
                                                     self.max_source_positions)
        self.register_buffer('token_rp_bucket', token_rp_bucket)
        self.token_rel_pos_table_list = nn.ModuleList()

        # Register image relative position embedding table
        self.image_bucket_size = image_bucket_size
        image_num_rel_dis = (2 * image_bucket_size -
                             1) * (2 * image_bucket_size - 1) + 3
        image_rp_bucket = make_image_bucket_position(image_bucket_size,
                                                     image_num_rel_dis)
        self.register_buffer('image_rp_bucket', image_rp_bucket)
        self.image_rel_pos_table_list = nn.ModuleList()

        # Build encoder layers
        self.layers = nn.ModuleList()
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, num_layers)]
        for index in range(self.num_layers):
            layer = OFAEncoderLayer(
                embedding_dim=embedding_dim,
                num_heads=num_heads,
                dropout_rate=dropout_rate,
                drop_path_rate=dpr[index],
                scale_factor=attn_scale_factor,
                pre_norm=pre_norm,
            )
            self.layers.append(layer)
            token_pos_table = nn.Embedding(token_num_rel_dis, self.num_heads)
            image_pos_table = nn.Embedding(image_num_rel_dis, self.num_heads)
            nn.init.constant_(token_pos_table.weight, 0.)
            nn.init.constant_(image_pos_table.weight, 0.)
            self.token_rel_pos_table_list.append(token_pos_table)
            self.image_rel_pos_table_list.append(image_pos_table)

        if pre_norm:
            self.final_ln = nn.LayerNorm(embedding_dim)
        else:
            self.final_ln = None

    main_input_name = 'input_ids'

    def forward(self,
                input_ids,
                images,
                images_mask,
                output_attentions=False,
                output_hidden_states=False,
                sample_patch_num=None):
        padding_mask = input_ids.eq(self.padding_idx)
        has_pads = padding_mask.any()
        token_embedding = self.embed_tokens(input_ids)
        token_embedding = self.embedding_scale * token_embedding

        # Embed the token position
        src_pos_idx = torch.arange(input_ids.size(-1), device=input_ids.device)
        src_pos_idx = src_pos_idx.expand(*input_ids.shape).contiguous()
        pos_embedding = self.embed_positions(src_pos_idx)

        # Embed the input tokens
        x = self.process_embedding(
            embedding=token_embedding,
            type_tokens=input_ids.new_zeros(token_embedding.shape[:2]),
            pos_embedding=pos_embedding,
            embedding_ln=self.embedding_ln,
        )
        pos_embedding = self.pos_ln(pos_embedding)

        # Embed the input images
        if images is not None:
            (image_tokens, image_padding_mask, image_position_ids,
             image_pos_embedding) = self.get_image_tokens(
                 images,
                 sample_patch_num,
                 images_mask,
             )
            image_embedding = self.image_proj(image_tokens)

            image_x = self.process_embedding(
                embedding=image_embedding,
                type_tokens=input_ids.new_ones(image_embedding.shape[:2]),
                pos_embedding=image_pos_embedding,
                embedding_ln=self.image_embedding_ln,
            )
            image_pos_embedding = self.image_pos_ln(image_pos_embedding)

            x = torch.cat([image_x, x], dim=1)
            padding_mask = torch.cat([image_padding_mask, padding_mask], dim=1)
            pos_embedding = torch.cat([image_pos_embedding, pos_embedding],
                                      dim=1)

        # account for padding while computing the representation
        if has_pads:
            x = x * (1 - padding_mask.unsqueeze(-1).type_as(x))

        # Decoupled position embedding
        B, L = pos_embedding.shape[:2]
        pos_q = self.pos_q_linear(pos_embedding).view(
            B, L, self.num_heads, -1).transpose(1, 2) * self.pos_scaling
        pos_k = self.pos_k_linear(pos_embedding).view(B, L, self.num_heads,
                                                      -1).transpose(1, 2)
        abs_pos_bias = torch.matmul(pos_q, pos_k.transpose(2, 3))

        all_hidden_states = [] if output_hidden_states else None
        all_attentions = [] if output_attentions else None

        for idx, layer in enumerate(self.layers):
            if output_hidden_states:
                all_hidden_states.append(x)

            self_attn_bias = abs_pos_bias.clone()
            # Add decoupled position embedding for input tokens.
            token_len = input_ids.size(1)
            rel_pos_bias = self.get_rel_pos_bias(input_ids, idx)
            self_attn_bias[:, :, -token_len:, -token_len:] += rel_pos_bias

            # Add decoupled position embedding for images
            if images is not None:
                token_len = image_tokens.size(1)
                rel_pos_bias = self.get_image_rel_pos_bias(
                    image_position_ids, idx)
                self_attn_bias[:, :, :token_len, :token_len] += rel_pos_bias

            if has_pads:
                attention_mask = _expand_mask(padding_mask, dtype=x.dtype)
            else:
                attention_mask = None

            out = layer(
                x,
                attention_mask=attention_mask,
                attn_bias=self_attn_bias,
                output_attentions=output_attentions)
            x = out[0]

            if output_attentions:
                all_attentions.append(out[1])

        if output_hidden_states:
            all_hidden_states.append(x)

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

        return OFAEncoderOutput(
            last_hidden_state=x,  # (B, L, C)
            padding_mask=padding_mask,  # (B, L)
            position_embedding=pos_embedding,  # (B, L, C)
            hidden_states=all_hidden_states,  # list of (B, L, C)
            attentions=all_attentions,  # list of (B, num_heads, L, head_dims)
        )

    def get_image_tokens(self, images, sample_patch_num, images_mask):
        image_embedding = self.embed_images(images)[-1]
        B, C, H, W = image_embedding.shape
        num_patches = H * W

        padding_mask = images.new_zeros((B, num_patches)).bool()
        position_col = torch.arange(W).unsqueeze(0)
        position_row = torch.arange(H).unsqueeze(1) * self.image_bucket_size
        position_idx = (position_col + position_row + 1).view(-1)
        position_idx = position_idx.to(images.device).expand(B, num_patches)

        # (B, C, H, W) -> (B, C, H*W) -> (B, H*W, C)
        image_embedding = image_embedding.flatten(2).transpose(1, 2)
        if sample_patch_num is not None:
            patch_orders = torch.stack([
                torch.randperm(num_patches)[:sample_patch_num]
                for _ in range(B)
            ])
            num_patches = sample_patch_num
            image_embedding = image_embedding.gather(
                dim=1, index=patch_orders.unsqueeze(2).expand(-1, -1, C))
            padding_mask = padding_mask.gather(1, patch_orders)
            position_idx = position_idx.gather(1, patch_orders)

        pos_embedding = self.embed_image_positions(position_idx)
        padding_mask[~images_mask] = True
        return image_embedding, padding_mask, position_idx, pos_embedding

    def process_embedding(self,
                          embedding,
                          pos_embedding=None,
                          type_tokens=None,
                          embedding_ln=None):
        if self.entangle_position_embedding and pos_embedding is not None:
            embedding += pos_embedding
        if self.embed_type is not None:
            embedding += self.embed_type(type_tokens)
        if embedding_ln is not None:
            embedding = embedding_ln(embedding)
        embedding = self.dropout(embedding)

        return embedding

    def get_rel_pos_bias(self, x, idx):
        seq_len = x.size(1)
        rp_bucket = self.token_rp_bucket[:seq_len, :seq_len]
        values = F.embedding(rp_bucket,
                             self.token_rel_pos_table_list[idx].weight)
        values = values.unsqueeze(0).expand(x.size(0), -1, -1, -1)
        values = values.permute([0, 3, 1, 2])
        return values.contiguous()

    def get_image_rel_pos_bias(self, image_position_ids, idx):
        bsz, seq_len = image_position_ids.shape
        rp_bucket_size = self.image_rp_bucket.size(1)

        rp_bucket = self.image_rp_bucket.unsqueeze(0).expand(
            bsz, rp_bucket_size, rp_bucket_size).gather(
                1, image_position_ids[:, :, None].expand(
                    bsz, seq_len, rp_bucket_size)).gather(
                        2, image_position_ids[:, None, :].expand(
                            bsz, seq_len, seq_len))
        values = F.embedding(rp_bucket,
                             self.image_rel_pos_table_list[idx].weight)
        values = values.permute(0, 3, 1, 2)
        return values


class OFADecoder(BaseModule):
    """The decoder module of OFA.

    Args:
        embed_tokens (nn.Embedding): The embedding module to embed the
            input tokens.
        num_layers (int): The number of decoder layers. Defaults to 6.
        num_heads (int): The number of heads of attention. Defaults to 12.
        dropout_rate (float): The prob of dropout for embedding and
            transformer layers. Defaults to 0.
        drop_path_rate (float): The prob of droppath for transformer layers.
            Defaults to 0.
        max_target_positions (int): The maximum length of the input tokens.
            Defaults to 1024.
        code_image_size (int): The resolution of the generated image in the
            image infilling task. Defaults to 128.
        token_bucket_size (int): The token bucket size, it's used as the
            maximum relative position index in relative position embedding
            of input tokens. Defaults to 256.
        image_bucket_size (int): The image bucket size, it's used to generate
            the image relative position embedding table. It should be larger
            than the shape of image feature map. Defaults to 42.
        attn_scale_factor (float): The scale factor to calculate qk scale in
            attentions. Defaults to 2.
        scale_embedding (bool): Whether to scale the embeddings by the square
            root of the dimension. Defaults to False.
        add_embedding_ln (bool): Whether to add an extra layer norm for token
            embeddings. Defaults to True.
        add_code_embedding_ln (bool): Whether to add an extra layer norm for
            code embeddings. Defaults to True.
        pre_norm (bool): Whether to do layer norm before attention and ffn
            blocks in transformer layers. Defaults to True.
        entangle_position_embedding (bool): Whether to add the position
            embedding on the embeddings directly. Defaults to False.
        share_input_output_embed (bool): Share the weights of the input token
            embedding module and the output projection module.
            Defaults to True.
        init_cfg (dict, optional): The initialization config. Defaults to None.
    """

    def __init__(
        self,
        embed_tokens,
        num_layers=6,
        num_heads=12,
        dropout_rate=0.,
        drop_layer_rate=0.,
        drop_path_rate=0.,
        max_target_positions=1024,
        code_image_size=128,
        token_bucket_size=256,
        image_bucket_size=42,
        attn_scale_factor=2.,
        scale_embedding=False,
        add_embedding_ln=True,
        add_code_embedding_ln=True,
        pre_norm=True,
        entangle_position_embedding=False,
        share_input_output_embed=True,
        init_cfg=None,
    ):
        super().__init__(init_cfg=init_cfg)
        self._future_mask = torch.empty(0)

        self.num_layers = num_layers
        embedding_dim = embed_tokens.embedding_dim
        self.embedding_dim = embedding_dim
        self.padding_idx = embed_tokens.padding_idx
        self.max_target_positions = max_target_positions
        self.num_heads = num_heads

        # Build embedding process components
        self.embed_tokens = embed_tokens
        self.embedding_scale = math.sqrt(
            embedding_dim) if scale_embedding else 1.0

        if add_embedding_ln:
            self.embedding_ln = nn.LayerNorm(embedding_dim)
        else:
            self.embedding_ln = None

        if add_code_embedding_ln:
            self.code_embedding_ln = nn.LayerNorm(embedding_dim)
        else:
            self.code_embedding_ln = None

        # Build position embedding
        self.embed_positions = nn.Embedding(self.max_target_positions + 2,
                                            embedding_dim)
        self.pos_ln = nn.LayerNorm(embedding_dim)
        self.embed_image_positions = nn.Embedding(image_bucket_size**2 + 1,
                                                  embedding_dim)
        self.image_pos_ln = nn.LayerNorm(embedding_dim)

        self.pos_scaling = float(embedding_dim / num_heads *
                                 attn_scale_factor)**-0.5
        self.self_pos_q_linear = nn.Linear(embedding_dim, embedding_dim)
        self.self_pos_k_linear = nn.Linear(embedding_dim, embedding_dim)
        self.cross_pos_q_linear = nn.Linear(embedding_dim, embedding_dim)
        self.cross_pos_k_linear = nn.Linear(embedding_dim, embedding_dim)

        self.entangle_position_embedding = entangle_position_embedding

        self.dropout = nn.Dropout(
            dropout_rate) if dropout_rate > 0. else nn.Identity()
        if drop_layer_rate > 0.:
            raise NotImplementedError

        # Register token relative position embedding table
        self.token_bucket_size = token_bucket_size
        token_num_rel_dis = 2 * token_bucket_size - 1
        token_rp_bucket = make_token_bucket_position(token_bucket_size)
        self.register_buffer('token_rp_bucket', token_rp_bucket)
        self.token_rel_pos_table_list = nn.ModuleList()

        # Register image relative position embedding table
        self.image_bucket_size = image_bucket_size
        image_num_rel_dis = (2 * image_bucket_size -
                             1) * (2 * image_bucket_size - 1) + 3
        image_rp_bucket = make_image_bucket_position(image_bucket_size,
                                                     image_num_rel_dis)
        self.register_buffer('image_rp_bucket', image_rp_bucket)
        self.image_rel_pos_table_list = nn.ModuleList()

        self.window_size = code_image_size // 8

        position_col = torch.arange(self.window_size).unsqueeze(0)
        position_row = torch.arange(
            self.window_size).unsqueeze(1) * self.image_bucket_size
        image_position_idx = (position_col + position_row + 1)
        image_position_idx = torch.cat(
            [torch.tensor([0]), image_position_idx.view(-1)])
        image_position_idx = torch.cat(
            [image_position_idx,
             torch.tensor([1024] * 768)])
        self.register_buffer('image_position_idx', image_position_idx)

        # Build decoder layers
        self.layers = nn.ModuleList()
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, num_layers)]
        for index in range(self.num_layers):
            layer = OFADecoderLayer(
                embedding_dim=embedding_dim,
                num_heads=num_heads,
                dropout_rate=dropout_rate,
                drop_path_rate=dpr[index],
                scale_factor=attn_scale_factor,
                pre_norm=pre_norm,
            )
            self.layers.append(layer)
            token_pos_table = nn.Embedding(token_num_rel_dis, self.num_heads)
            image_pos_table = nn.Embedding(image_num_rel_dis, self.num_heads)
            nn.init.constant_(token_pos_table.weight, 0.)
            nn.init.constant_(image_pos_table.weight, 0.)
            self.token_rel_pos_table_list.append(token_pos_table)
            self.image_rel_pos_table_list.append(image_pos_table)

        if pre_norm:
            self.final_ln = nn.LayerNorm(embedding_dim)
        else:
            self.final_ln = None

        # Build output projection
        if share_input_output_embed:
            self.output_projection = nn.Linear(
                self.embed_tokens.weight.shape[1],
                self.embed_tokens.weight.shape[0],
                bias=False,
            )
            self.output_projection.weight = self.embed_tokens.weight
        else:
            vocab_size = self.embed_tokens.num_embeddings
            self.output_projection = nn.Linear(
                embedding_dim, vocab_size, bias=False)
            nn.init.normal_(
                self.output_projection.weight,
                mean=0,
                std=embedding_dim**-0.5,
            )

    main_input_name = 'input_ids'

    def forward(
        self,
        input_ids: torch.Tensor = None,
        attention_mask: torch.Tensor = None,
        encoder_hidden_states: torch.Tensor = None,
        encoder_attention_mask: torch.Tensor = None,
        code_masks: Optional[torch.Tensor] = None,
        encoder_pos_embedding: Optional[torch.Tensor] = None,
        past_key_values: Optional[torch.Tensor] = None,
        use_cache: bool = False,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
    ):

        if past_key_values is not None and len(past_key_values) > 0:
            B, _, L_past, _ = past_key_values[0][0].shape
            L = L_past + 1
        else:
            B, L = input_ids.shape
            L_past = 0

        # Embed the token position
        target_pos_idx = torch.arange(
            L, device=input_ids.device).expand([B, L]).contiguous()
        pos_embedding = self.embed_positions(target_pos_idx)

        # Embed the code positions
        if code_masks is not None and torch.any(code_masks):
            image_position_idx = self.image_position_idx[:input_ids.size(1)]
            image_position_idx = image_position_idx.unsqueeze(0).expand(B, L)
            pos_embedding[code_masks] = self.embed_image_positions(
                image_position_idx)[code_masks]

        # Self-attention position bias (B, num_heads, L_t, L_t)
        self_abs_pos_bias = self.get_pos_info(self.pos_ln(pos_embedding))
        if code_masks is not None and torch.any(code_masks):
            self_image_abs_pos_bias = self.get_pos_info(
                self.image_pos_ln(pos_embedding))
            self_abs_pos_bias[code_masks] = self_image_abs_pos_bias[code_masks]

        # Cross-attention position bias (B, num_heads, L_t, L_s)
        cross_abs_pos_bias = self.get_pos_info(
            self.pos_ln(pos_embedding), encoder_pos_embedding)
        if code_masks is not None and torch.any(code_masks):
            cross_image_abs_pos_bias = self.get_pos_info(
                self.image_pos_ln(pos_embedding), encoder_pos_embedding)
            cross_abs_pos_bias[code_masks] = cross_image_abs_pos_bias[
                code_masks]

        all_prev_output_tokens = input_ids.clone()
        if past_key_values is not None and len(past_key_values) > 0:
            input_ids = input_ids[:, -1:]
            cross_abs_pos_bias = cross_abs_pos_bias[:, :, -1:, :]
            pos_embedding = pos_embedding[:, -1:, :]

        # Embed the input tokens
        x = self.embed_tokens(input_ids) * self.embedding_scale

        if self.entangle_position_embedding:
            x += pos_embedding

        if self.embedding_ln is not None:
            if (code_masks is None or not code_masks.any()
                    or self.code_embedding_ln is None):
                x = self.embedding_ln(x)
            elif code_masks is not None and code_masks.all():
                x = self.code_embedding_ln(x)
            else:
                x[~code_masks] = self.embedding_ln(x[~code_masks])
                x[code_masks] = self.code_embedding_ln(x[code_masks])

        x = self.dropout(x)

        attention_mask = self._prepare_decoder_attention_mask(
            attention_mask, input_ids.shape, x.dtype, L_past)
        attention_mask = attention_mask.to(x.device)

        # decoder layers
        all_hidden_states = [] if output_hidden_states else None
        all_self_attns = [] if output_attentions else None
        all_cross_attentions = [] if (
            output_attentions and encoder_hidden_states is not None) else None
        next_decoder_cache = [] if use_cache else None

        for idx, layer in enumerate(self.layers):
            # add hidden states from the last decoder layer
            if output_hidden_states:
                all_hidden_states.append(x)

            if past_key_values is not None and len(past_key_values) > 0:
                past_key_value = past_key_values[idx]
            else:
                past_key_value = None

            self_attn_bias = self_abs_pos_bias.clone()
            if code_masks is None or not code_masks.any():
                self_attn_bias += self.get_rel_pos_bias(
                    all_prev_output_tokens, idx)
            elif code_masks is not None and code_masks.all():
                self_attn_bias += self.get_image_rel_pos_bias(
                    all_prev_output_tokens, idx)
            else:
                self_attn_bias[~code_masks] += self.get_rel_pos_bias(
                    all_prev_output_tokens, idx)
                self_attn_bias[code_masks] += self.get_image_rel_pos_bias(
                    all_prev_output_tokens, idx)

            if past_key_value is not None:
                self_attn_bias = self_attn_bias[:, :, -1:, :]

            out = layer(
                x,
                attention_mask=attention_mask,
                encoder_hidden_states=encoder_hidden_states,
                encoder_attention_mask=encoder_attention_mask,
                past_key_value=past_key_value,
                output_attentions=output_attentions,
                use_cache=use_cache,
                self_attn_bias=self_attn_bias,
                cross_attn_bias=cross_abs_pos_bias,
            )
            x = out.pop(0)

            if output_attentions:
                all_self_attns.append(out.pop(0))
                if encoder_hidden_states is not None:
                    all_cross_attentions.append(out.pop(0))

            if use_cache:
                next_decoder_cache.append(out.pop(0))

        # add hidden states from the last decoder layer
        if output_hidden_states:
            all_hidden_states += (x, )

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

        x = self.output_projection(x)

        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=x,
            past_key_values=next_decoder_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
            cross_attentions=all_cross_attentions,
        )

    def _prepare_decoder_attention_mask(
        self,
        attention_mask,
        input_shape,
        dtype,
        past_key_values_length,
    ):
        r"""
        Create causal mask for unidirectional decoding.
        [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
        """
        combined_attention_mask = None
        if input_shape[-1] > 1:
            combined_attention_mask = _make_causal_mask(
                input_shape,
                dtype,
                past_key_values_length=past_key_values_length).to(
                    attention_mask.device)

        if attention_mask is not None:
            # (B, L_s) -> (B, 1, L_t, L_s)
            expanded_attention_mask = _expand_mask(
                attention_mask, dtype, tgt_len=input_shape[-1])
            combined_attention_mask = (
                expanded_attention_mask if combined_attention_mask is None else
                expanded_attention_mask + combined_attention_mask)

        return combined_attention_mask

    def get_pos_info(self, pos_embedding, src_pos_embedding=None):
        B, tgt_len = pos_embedding.shape[:2]
        if src_pos_embedding is not None:
            src_len = src_pos_embedding.size(1)
            pos_q = self.cross_pos_q_linear(pos_embedding).view(
                B, tgt_len, self.num_heads, -1).transpose(1, 2)
            pos_q = pos_q * self.pos_scaling
            pos_k = self.cross_pos_k_linear(src_pos_embedding).view(
                B, src_len, self.num_heads, -1).transpose(1, 2)
        else:
            pos_q = self.self_pos_q_linear(pos_embedding).view(
                B, tgt_len, self.num_heads, -1).transpose(1, 2)
            pos_q = pos_q * self.pos_scaling
            pos_k = self.self_pos_k_linear(pos_embedding).view(
                B, tgt_len, self.num_heads, -1).transpose(1, 2)

        abs_pos_bias = torch.matmul(pos_q, pos_k.transpose(2, 3))

        return abs_pos_bias

    def get_rel_pos_bias(self, x, idx):
        seq_len = x.size(1)
        rp_bucket = self.token_rp_bucket[:seq_len, :seq_len]
        values = F.embedding(rp_bucket,
                             self.token_rel_pos_table_list[idx].weight)
        values = values.unsqueeze(0).expand(x.size(0), -1, -1, -1)
        values = values.permute([0, 3, 1, 2])
        return values.contiguous()

    def get_image_rel_pos_bias(self, image_position_ids, idx):
        bsz, seq_len = image_position_ids.shape
        rp_bucket_size = self.image_rp_bucket.size(1)

        rp_bucket = self.image_rp_bucket.unsqueeze(0).expand(
            bsz, rp_bucket_size, rp_bucket_size).gather(
                1, image_position_ids[:, :, None].expand(
                    bsz, seq_len, rp_bucket_size)).gather(
                        2, image_position_ids[:, None, :].expand(
                            bsz, seq_len, seq_len))
        values = F.embedding(rp_bucket,
                             self.image_rel_pos_table_list[idx].weight)
        values = values.permute(0, 3, 1, 2)
        return values


class OFAEncoderDecoder(BaseModule, GenerationMixin):
    """The OFA main architecture with an encoder and a decoder.

    Args:
        encoder_cfg (dict): The config of the encoder, accept the keyword
            arguments of :class:`OFAEncoder`.
        decoder_cfg (dict): The config of the decoder, accept the keyword
            arguments of :class:`OFADecoder`.
        padding_idx (int): The index of the padding token.
        vocab_size (int): The size of the vocabulary.
        embedding_dim (int): The embedding dimensions of both the encoder
            and the decoder.
        generation_cfg (dict): The extra generation config, accept the keyword
            arguments of :class:`~transformers.GenerationConfig`.
            Defaults to an empty dict.
        init_cfg (dict, optional): The initialization config. Defaults to None.
    """

    def __init__(
            self,
            encoder_cfg,
            decoder_cfg,
            padding_idx,
            vocab_size,
            embedding_dim,
            generation_cfg=dict(),
            init_cfg=None,
    ):
        super().__init__(init_cfg=init_cfg)

        self.padding_idx = padding_idx
        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim
        embed_tokens = nn.Embedding(vocab_size, embedding_dim, padding_idx)

        self.encoder = OFAEncoder(embed_tokens, **encoder_cfg)
        self.decoder = OFADecoder(embed_tokens, **decoder_cfg)

        self.config = PretrainedConfig(
            vocab_size=vocab_size,
            embedding_dim=embedding_dim,
            padding_idx=padding_idx,
            bos_token_id=0,
            decoder_start_token_id=0,
            pad_token_id=1,
            eos_token_id=2,
            forced_eos_token_id=2,
            use_cache=False,
            is_encoder_decoder=True,
        )
        self.config.update(generation_cfg)

        self.generation_config = GenerationConfig.from_model_config(
            self.config)

    @property
    def device(self):
        return next(self.parameters()).device

    def can_generate(self):
        return True

    def get_encoder(self):
        return self.encoder

    def get_decoder(self):
        return self.decoder

    def max_decoder_positions(self):
        """Maximum length supported by the decoder."""
        return self.decoder.max_positions()

    def get_normalized_probs(self, net_output, log_probs: bool, sample=None):
        """Get normalized probabilities (or log probs) from a net's output."""
        return self.get_normalized_probs_scriptable(net_output, log_probs,
                                                    sample)

    def get_normalized_probs_scriptable(
        self,
        net_output,
        log_probs: bool,
        sample=None,
    ):
        """Scriptable helper function for get_normalized_probs in.

        ~BaseFairseqModel.
        """
        if hasattr(self, 'decoder'):
            return self.decoder.get_normalized_probs(net_output, log_probs,
                                                     sample)
        elif torch.is_tensor(net_output):
            # syntactic sugar for simple models which don't have a decoder
            # (e.g., the classification tutorial)
            logits = net_output.float()
            if log_probs:
                return F.log_softmax(logits, dim=-1)
            else:
                return F.softmax(logits, dim=-1)
        raise NotImplementedError

    main_input_name = 'input_ids'

    def forward(self,
                input_ids=None,
                images=None,
                images_mask=None,
                sample_patch_num=None,
                decoder_input_ids=None,
                code_masks=None,
                attention_mask=None,
                encoder_outputs=None,
                past_key_values=None,
                use_cache=False,
                output_attentions=False,
                output_hidden_states=False,
                constrain_fn=None,
                return_dict=False):
        """Forword the module.

        Args:
            input_ids (torch.Tensor): The indices of the input tokens in the
                vocabulary, and padding will be ignored by default. The indices
                can be obtained using :class:`OFATokenizer`.
                The shape is (B, L).
            images (torch.Tensor): The input images. The shape is (B, 3, H, W).
            images_mask (torch.Tensor): The mask of all available images. The
                shape is (B, ).
            sample_patch_num (int): The number of patches to sample for the
                images. Defaults to None, which means to use all patches.
            decoder_input_ids (torch.Tensor): The indices of the input tokens
                for the decoder.
            code_masks (torch.Tensor): The mask of all samples for image
                generation. The shape is (B, ).
            attention_mask (torch.Tensor): The attention mask for decoding.
                The shape is (B, L).
            encoder_outputs (OFAEncoderOutput): The encoder outputs with hidden
                states, positional embeddings, and padding masks.
            past_key_values (Tuple[Tuple[torch.Tensor]]): If use cache, the
                parameter is a tuple of length ``num_layers``. Every item is
                also a tuple with four tensors, two for the key and value of
                self-attention, two for the key and value of cross-attention.
            use_cache (bool): Whether to use cache for faster inference.
                Defaults to False.
            output_attentions (bool): Whether to output attention weights.
                Defaults to False.
            output_hidden_states (bool): Whether to output hidden states.
                Defaults to False.
            constrain_fn (Callable, optional): The function to constrain the
                output logits. Defaults to None.
            return_dict (bool): Not used, it's only for compat with the
                interface of the ``generate`` of ``transformers``.

        Returns:
            Seq2SeqLMOutput:

            - logits (``torch.Tensor``): The last decoder hidden states.
              The shape is (B, L, C).
            - past_key_values (``Tuple[Tuple[torch.Tensor]]``): The past keys
              and values for faster inference.
            - decoder_hidden_states (``Tuple[torch.Tensor]``): the decoder
              hidden states of all layers.
            - decoder_attentions (``Tuple[torch.Tensor]``): The self-attention
              weights of all layers in the decoder.
            - cross_attentions (``Tuple[torch.Tensor]``): The cross-attention
              weights of all layers in the decoder.
            - encoder_last_hidden_state (``torch.Tensor``): The last encoder
              hidden states.
            - encoder_hidden_states (``Tuple[torch.Tensor]``): The encoder
              hidden states of all layers, including the embeddings.
            - encoder_attentions (``Tuple[torch.Tensor]``): The self-attention
              weights of all layers in the encoder.
        """

        if encoder_outputs is None:
            encoder_outputs = self.encoder(
                input_ids=input_ids,
                images=images,
                images_mask=images_mask,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                sample_patch_num=sample_patch_num,
            )

        if decoder_input_ids.eq(self.padding_idx).any():
            attention_mask = decoder_input_ids.eq(self.padding_idx)

        encoder_hidden_states = encoder_outputs.last_hidden_state
        encoder_attention_mask = _expand_mask(encoder_outputs.padding_mask,
                                              encoder_hidden_states.dtype,
                                              decoder_input_ids.shape[-1])
        src_pos_embed = encoder_outputs.position_embedding

        decoder_outputs = self.decoder(
            input_ids=decoder_input_ids,
            attention_mask=attention_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            code_masks=code_masks,
            encoder_pos_embedding=src_pos_embed,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
        )

        # The constrain operation for fine-tuned model in OFA is applied
        # before log_softmax, therefore we cannot use
        # `prefix_allowed_tokens_fn` to implement it.
        if constrain_fn is not None:
            logits = constrain_fn(decoder_input_ids,
                                  decoder_outputs.last_hidden_state)
        else:
            logits = decoder_outputs.last_hidden_state

        return Seq2SeqLMOutput(
            logits=logits,
            past_key_values=decoder_outputs.past_key_values,
            decoder_hidden_states=decoder_outputs.hidden_states,
            decoder_attentions=decoder_outputs.attentions,
            cross_attentions=decoder_outputs.cross_attentions,
            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
            encoder_hidden_states=encoder_outputs.hidden_states,
            encoder_attentions=encoder_outputs.attentions,
        )

    def prepare_inputs_for_generation(self,
                                      decoder_input_ids=None,
                                      past=None,
                                      attention_mask=None,
                                      code_masks=None,
                                      use_cache=False,
                                      encoder_outputs=None,
                                      constrain_fn=None,
                                      **kwargs):
        # if attention_mask is None:
        attention_mask = decoder_input_ids.new_zeros(decoder_input_ids.shape)

        # cut decoder_input_ids if past is used
        if past is not None:
            decoder_input_ids = decoder_input_ids[:, -1:]

        return {
            'input_ids': None,
            'images': None,
            'images_mask': None,
            'sample_patch_num': None,
            'attention_mask': attention_mask,
            'encoder_outputs': encoder_outputs,
            'past_key_values': past,
            'decoder_input_ids': decoder_input_ids,
            'code_masks': code_masks,
            'use_cache': use_cache,
            'constrain_fn': constrain_fn,
        }

    def _prepare_encoder_decoder_kwargs_for_generation(
            self,
            inputs_tensor: torch.Tensor,
            model_kwargs,
            model_input_name: Optional[str] = None):
        # 1. get encoder
        encoder = self.get_encoder()

        # 2. prepare encoder args and encoder kwargs from model kwargs
        irrelevant_prefix = [
            'decoder_', 'cross_attn', 'use_cache', 'attention_mask',
            'constrain_fn'
        ]
        encoder_kwargs = {
            argument: value
            for argument, value in model_kwargs.items()
            if not any(argument.startswith(p) for p in irrelevant_prefix)
        }

        if encoder_kwargs.get('images_mask') is None:
            encoder_kwargs['images_mask'] = torch.tensor([True] *
                                                         inputs_tensor.size(0))

        # 3. make sure that encoder returns `ModelOutput`
        model_input_name = model_input_name or self.main_input_name
        encoder_kwargs[model_input_name] = inputs_tensor
        model_kwargs['encoder_outputs']: ModelOutput = encoder(
            **encoder_kwargs)
        model_kwargs['attention_mask'] = None

        return model_kwargs

    @staticmethod
    def _reorder_cache(past, beam_idx):
        reordered_past = ()
        for layer_past in past:
            reordered_past += (tuple(
                past_state.index_select(0, beam_idx)
                for past_state in layer_past), )
        return reordered_past

    @staticmethod
    def _expand_inputs_for_generation(
        input_ids: torch.LongTensor,
        expand_size: int = 1,
        is_encoder_decoder: bool = False,
        attention_mask: Optional[torch.LongTensor] = None,
        encoder_outputs: Optional[ModelOutput] = None,
        **model_kwargs,
    ):
        expanded_return_idx = (
            torch.arange(input_ids.shape[0]).view(-1, 1).repeat(
                1, expand_size).view(-1).to(input_ids.device))
        input_ids = input_ids.index_select(0, expanded_return_idx)

        if attention_mask is not None:
            model_kwargs['attention_mask'] = attention_mask.index_select(
                0, expanded_return_idx)

        if is_encoder_decoder:
            if encoder_outputs is None:
                raise ValueError('If `is_encoder_decoder` is True, make '
                                 'sure that `encoder_outputs` is defined.')
            encoder_outputs['last_hidden_state'] = encoder_outputs.\
                last_hidden_state.index_select(0, expanded_return_idx)
            encoder_outputs['position_embedding'] = encoder_outputs.\
                position_embedding.index_select(0, expanded_return_idx)
            encoder_outputs['padding_mask'] = encoder_outputs.\
                padding_mask.index_select(0, expanded_return_idx)
            model_kwargs['encoder_outputs'] = encoder_outputs
        return input_ids, model_kwargs