AutoAndroidController/py/venv/Lib/site-packages/modelscope/models/nlp/unite/translation_evaluation.py

# Copyright (c) Alibaba, Inc. and its affiliates.
"""PyTorch UniTE model."""

import warnings
from dataclasses import dataclass
from math import ceil
from typing import Dict, List, Optional, Tuple, Union

import numpy as np
import torch
import torch.utils.checkpoint
from packaging import version
from torch.nn import (Dropout, Linear, Module, Parameter, ParameterList,
                      Sequential)
from torch.nn.functional import softmax
from torch.nn.utils.rnn import pad_sequence
from transformers import XLMRobertaConfig, XLMRobertaModel
from transformers.activations import ACT2FN

from modelscope.metainfo import Models
from modelscope.models.base import TorchModel
from modelscope.models.builder import MODELS
from modelscope.models.nlp.unite.configuration import InputFormat
from modelscope.outputs.nlp_outputs import TranslationEvaluationOutput
from modelscope.utils.compatible_with_transformers import \
    compatible_position_ids
from modelscope.utils.constant import Tasks
from modelscope.utils.logger import get_logger

logger = get_logger()

__all__ = ['UniTEForTranslationEvaluation']


def _layer_norm_all(tensor, mask_float):
    broadcast_mask = mask_float.unsqueeze(dim=-1)
    num_elements_not_masked = broadcast_mask.sum() * tensor.size(-1)
    tensor_masked = tensor * broadcast_mask

    mean = tensor_masked.sum([-1, -2, -3],
                             keepdim=True) / num_elements_not_masked
    variance = (((tensor_masked - mean) * broadcast_mask)**2).sum(
        [-1, -2, -3], keepdim=True) / num_elements_not_masked

    return (tensor - mean) / torch.sqrt(variance + 1e-12)


class LayerwiseAttention(Module):

    def __init__(
        self,
        num_layers: int,
        model_dim: int,
        dropout: float = None,
    ) -> None:
        super(LayerwiseAttention, self).__init__()
        self.num_layers = num_layers
        self.model_dim = model_dim
        self.dropout = dropout

        self.scalar_parameters = Parameter(
            torch.zeros((num_layers, ), requires_grad=True))
        self.gamma = Parameter(torch.FloatTensor([1.0]), requires_grad=True)

        if self.dropout:
            dropout_mask = torch.zeros(len(self.scalar_parameters))
            dropout_fill = torch.empty(len(
                self.scalar_parameters)).fill_(-1e20)
            self.register_buffer('dropout_mask', dropout_mask)
            self.register_buffer('dropout_fill', dropout_fill)

    def forward(
        self,
        tensors: List[torch.Tensor],  # pylint: disable=arguments-differ
        mask: torch.Tensor = None,
    ) -> torch.Tensor:
        tensors = torch.cat(list(x.unsqueeze(dim=0) for x in tensors), dim=0)

        if self.training and self.dropout:
            normed_weights = softmax(
                torch.where(self.dropout_mask.uniform_() > self.dropout,
                            self.scalar_parameters, self.dropout_fill),
                dim=-1)
        else:
            normed_weights = softmax(self.scalar_parameters, dim=-1)

        normed_weights = normed_weights.view(-1, 1, 1, 1)

        mask_float = mask.float()
        weighted_sum = (normed_weights
                        * _layer_norm_all(tensors, mask_float)).sum(dim=0)
        weighted_sum = weighted_sum[:, 0, :]

        return self.gamma * weighted_sum


class FeedForward(Module):

    def __init__(
        self,
        in_dim: int,
        out_dim: int = 1,
        hidden_sizes: List[int] = [3072, 768],
        activations: str = 'Sigmoid',
        final_activation: Optional[str] = None,
        dropout: float = 0.1,
    ) -> None:
        """
        Feed Forward Neural Network.

        Args:
            in_dim (:obj:`int`):
                Number of input features.
            out_dim (:obj:`int`, defaults to 1):
                Number of output features. Default is 1 -- a single scalar.
            hidden_sizes (:obj:`List[int]`, defaults to `[3072, 768]`):
                List with hidden layer sizes.
            activations (:obj:`str`, defaults to `Sigmoid`):
                Name of the activation function to be used in the hidden layers.
            final_activation (:obj:`str`, Optional, defaults to `None`):
                Name of the final activation function if any.
            dropout (:obj:`float`, defaults to 0.1):
                Dropout ratio to be used in the hidden layers.
        """
        super().__init__()
        modules = []
        modules.append(Linear(in_dim, hidden_sizes[0]))
        modules.append(self.build_activation(activations))
        modules.append(Dropout(dropout))

        for i in range(1, len(hidden_sizes)):
            modules.append(Linear(hidden_sizes[i - 1], hidden_sizes[i]))
            modules.append(self.build_activation(activations))
            modules.append(Dropout(dropout))

        modules.append(Linear(hidden_sizes[-1], int(out_dim)))
        if final_activation is not None:
            modules.append(self.build_activation(final_activation))

        self.ff = Sequential(*modules)

    def build_activation(self, activation: str) -> Module:
        return ACT2FN[activation]

    def forward(self, in_features: torch.Tensor) -> torch.Tensor:
        return self.ff(in_features)


@MODELS.register_module(Tasks.translation_evaluation, module_name=Models.unite)
class UniTEForTranslationEvaluation(TorchModel):

    def __init__(self,
                 attention_probs_dropout_prob: float = 0.1,
                 bos_token_id: int = 0,
                 eos_token_id: int = 2,
                 pad_token_id: int = 1,
                 hidden_act: str = 'gelu',
                 hidden_dropout_prob: float = 0.1,
                 hidden_size: int = 1024,
                 initializer_range: float = 0.02,
                 intermediate_size: int = 4096,
                 layer_norm_eps: float = 1e-05,
                 max_position_embeddings: int = 512,
                 num_attention_heads: int = 16,
                 num_hidden_layers: int = 24,
                 type_vocab_size: int = 1,
                 use_cache: bool = True,
                 vocab_size: int = 250002,
                 mlp_hidden_sizes: List[int] = [3072, 1024],
                 mlp_act: str = 'tanh',
                 mlp_final_act: Optional[str] = None,
                 mlp_dropout: float = 0.1,
                 **kwargs):
        r"""The UniTE Model which outputs the scalar to describe the corresponding
            translation quality of hypothesis. The model architecture includes two
            modules: a pre-trained language model (PLM) to derive representations,
            and a multi-layer perceptron (MLP) to give predicted score.

            Args:
                attention_probs_dropout_prob (:obj:`float`, defaults to 0.1):
                    The dropout ratio for attention weights inside PLM.
                bos_token_id (:obj:`int`, defaults to 0):
                    The numeric id representing beginning-of-sentence symbol.
                eos_token_id (:obj:`int`, defaults to 2):
                    The numeric id representing ending-of-sentence symbol.
                pad_token_id (:obj:`int`, defaults to 1):
                    The numeric id representing padding symbol.
                hidden_act (:obj:`str`, defaults to :obj:`"gelu"`):
                    Activation inside PLM.
                hidden_dropout_prob (:obj:`float`, defaults to 0.1):
                    The dropout ratio for activation states inside PLM.
                hidden_size (:obj:`int`, defaults to 1024):
                    The dimensionality of PLM.
                initializer_range (:obj:`float`, defaults to 0.02):
                    The hyper-parameter for initializing PLM.
                intermediate_size (:obj:`int`, defaults to 4096):
                    The dimensionality of PLM inside feed-forward block.
                layer_norm_eps (:obj:`float`, defaults to 1e-5):
                    The value for setting epsilon to avoid zero-division inside
                        layer normalization.
                max_position_embeddings: (:obj:`int`, defaults to 512):
                    The maximum value for identifying the length of input sequence.
                num_attention_heads (:obj:`int`, defaults to 16):
                    The number of attention heads inside multi-head attention layer.
                num_hidden_layers (:obj:`int`, defaults to 24):
                    The number of layers inside PLM.
                type_vocab_size (:obj:`int`, defaults to 1):
                    The number of type embeddings.
                use_cache (:obj:`bool`, defaults to :obj:`True`):
                    Whether to use cached buffer to initialize PLM.
                vocab_size (:obj:`int`, defaults to 250002):
                    The size of vocabulary.
                mlp_hidden_sizes (:obj:`List[int]`, defaults to `[3072, 1024]`):
                    The size of hidden states inside MLP.
                mlp_act (:obj:`str`, defaults to :obj:`"tanh"`):
                    Activation inside MLP.
                mlp_final_act (:obj:`str`, `optional`, defaults to :obj:`None`):
                    Activation at the end of MLP.
                mlp_dropout (:obj:`float`, defaults to 0.1):
                    The dropout ratio for MLP.
            """
        super().__init__(**kwargs)

        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.bos_token_id = bos_token_id
        self.eos_token_id = eos_token_id
        self.pad_token_id = pad_token_id
        self.hidden_act = hidden_act
        self.hidden_dropout_prob = hidden_dropout_prob
        self.hidden_size = hidden_size
        self.initializer_range = initializer_range
        self.intermediate_size = intermediate_size
        self.layer_norm_eps = layer_norm_eps
        self.max_position_embeddings = max_position_embeddings
        self.num_attention_heads = num_attention_heads
        self.num_hidden_layers = num_hidden_layers
        self.type_vocab_size = type_vocab_size
        self.use_cache = use_cache
        self.vocab_size = vocab_size
        self.mlp_hidden_sizes = mlp_hidden_sizes
        self.mlp_act = mlp_act
        self.mlp_final_act = mlp_final_act
        self.mlp_dropout = mlp_dropout

        self.encoder_config = XLMRobertaConfig(
            bos_token_id=self.bos_token_id,
            eos_token_id=self.eos_token_id,
            pad_token_id=self.pad_token_id,
            vocab_size=self.vocab_size,
            hidden_size=self.hidden_size,
            num_hidden_layers=self.num_hidden_layers,
            num_attention_heads=self.num_attention_heads,
            intermediate_size=self.intermediate_size,
            hidden_act=self.hidden_act,
            hidden_dropout_prob=self.hidden_dropout_prob,
            attention_probs_dropout_prob=self.attention_probs_dropout_prob,
            max_position_embeddings=self.max_position_embeddings,
            type_vocab_size=self.type_vocab_size,
            initializer_range=self.initializer_range,
            layer_norm_eps=self.layer_norm_eps,
            use_cache=self.use_cache)

        self.encoder = XLMRobertaModel(
            self.encoder_config, add_pooling_layer=False)

        self.layerwise_attention = LayerwiseAttention(
            num_layers=self.num_hidden_layers + 1,
            model_dim=self.hidden_size,
            dropout=self.mlp_dropout)

        self.estimator = FeedForward(
            in_dim=self.hidden_size,
            out_dim=1,
            hidden_sizes=self.mlp_hidden_sizes,
            activations=self.mlp_act,
            final_activation=self.mlp_final_act,
            dropout=self.mlp_dropout)

        return

    def forward(self,
                input_ids: torch.Tensor,
                input_format: Optional[List[InputFormat]] = None,
                score: Optional[torch.Tensor] = None,
                **kwargs) -> TranslationEvaluationOutput:
        attention_mask = input_ids.ne(self.pad_token_id).long()
        outputs = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True,
            return_dict=True)
        mix_states = self.layerwise_attention(outputs['hidden_states'],
                                              attention_mask)
        pred = self.estimator(mix_states).squeeze(dim=-1)
        output = TranslationEvaluationOutput(
            score=pred.cpu().tolist(), input_format=input_format)

        if score is not None:
            loss = (pred - score).pow(2).mean()
            output['loss'] = loss

        return output

    def load_checkpoint(self, path: str, device: torch.device, plm_only: bool):
        if plm_only:
            self.encoder = self.encoder.from_pretrained(path).to(device)
            self.encoder.pooler = None
        else:
            state_dict = torch.load(path, map_location=device)
            compatible_position_ids(state_dict,
                                    'encoder.embeddings.position_ids')
            self.load_state_dict(state_dict)
        logger.info('Loading checkpoint parameters from %s' % path)
        return


def combine_input_sentences(all_input_concat: List[List[torch.Tensor]],
                            maximum_length: int = 512,
                            pad_idx: int = 1,
                            eos_idx: int = 2):
    for group in all_input_concat[1:]:
        group[:, 0] = eos_idx

    if len(all_input_concat) == 3:
        return cut_long_sequences3(all_input_concat, maximum_length, pad_idx)
    else:
        return cut_long_sequences2(all_input_concat, maximum_length, pad_idx)


def cut_long_sequences2(all_input_concat: List[List[torch.Tensor]],
                        maximum_length: int = 512,
                        pad_idx: int = 1):
    all_input_concat = list(zip(*all_input_concat))
    collected_tuples = list()
    for tensor_tuple in all_input_concat:
        tensor_tuple = tuple(
            x.masked_select(x.ne(pad_idx)) for x in tensor_tuple)
        all_lens = tuple(len(x) for x in tensor_tuple)

        if sum(all_lens) > maximum_length:
            lengths = dict(enumerate(all_lens))
            lengths_sorted_idxes = list(x[0] for x in sorted(
                lengths.items(), key=lambda d: d[1], reverse=True))

            offset = ceil((sum(lengths.values()) - maximum_length) / 2)

            if min(all_lens) > (maximum_length
                                // 2) and min(all_lens) > offset:
                lengths = dict((k, v - offset) for k, v in lengths.items())
            else:
                lengths[lengths_sorted_idxes[0]] = maximum_length - lengths[
                    lengths_sorted_idxes[1]]

            new_lens = list(lengths[k] for k in range(0, len(tensor_tuple)))
            new_tensor_tuple = tuple(x[:y]
                                     for x, y in zip(tensor_tuple, new_lens))
            for x, y in zip(new_tensor_tuple, tensor_tuple):
                x[-1] = y[-1]
            collected_tuples.append(new_tensor_tuple)
        else:
            collected_tuples.append(tensor_tuple)

    concat_tensor = list(torch.cat(x, dim=0) for x in collected_tuples)
    all_input_concat_padded = pad_sequence(
        concat_tensor, batch_first=True, padding_value=pad_idx)
    return all_input_concat_padded


def cut_long_sequences3(all_input_concat: List[List[torch.Tensor]],
                        maximum_length: int = 512,
                        pad_idx: int = 1):
    all_input_concat = list(zip(*all_input_concat))
    collected_tuples = list()
    for tensor_tuple in all_input_concat:
        tensor_tuple = tuple(
            x.masked_select(x.ne(pad_idx)) for x in tensor_tuple)
        all_lens = tuple(len(x) for x in tensor_tuple)

        if sum(all_lens) > maximum_length:
            lengths = dict(enumerate(all_lens))
            lengths_sorted_idxes = list(x[0] for x in sorted(
                lengths.items(), key=lambda d: d[1], reverse=True))

            offset = ceil((sum(lengths.values()) - maximum_length) / 3)

            if min(all_lens) > (maximum_length
                                // 3) and min(all_lens) > offset:
                lengths = dict((k, v - offset) for k, v in lengths.items())
            else:
                while sum(lengths.values()) > maximum_length:
                    if lengths[lengths_sorted_idxes[0]] > lengths[
                            lengths_sorted_idxes[1]]:
                        offset = maximum_length - lengths[lengths_sorted_idxes[
                            1]] - lengths[lengths_sorted_idxes[2]]
                        if offset > lengths[lengths_sorted_idxes[1]]:
                            lengths[lengths_sorted_idxes[0]] = offset
                        else:
                            lengths[lengths_sorted_idxes[0]] = lengths[
                                lengths_sorted_idxes[1]]
                    elif lengths[lengths_sorted_idxes[0]] == lengths[
                            lengths_sorted_idxes[1]] > lengths[
                                lengths_sorted_idxes[2]]:
                        offset = (maximum_length
                                  - lengths[lengths_sorted_idxes[2]]) // 2
                        if offset > lengths[lengths_sorted_idxes[2]]:
                            lengths[lengths_sorted_idxes[0]] = lengths[
                                lengths_sorted_idxes[1]] = offset
                        else:
                            lengths[lengths_sorted_idxes[0]] = lengths[
                                lengths_sorted_idxes[1]] = lengths[
                                    lengths_sorted_idxes[2]]
                    else:
                        lengths[lengths_sorted_idxes[0]] = lengths[
                            lengths_sorted_idxes[1]] = lengths[
                                lengths_sorted_idxes[2]] = maximum_length // 3

            new_lens = list(lengths[k] for k in range(0, len(lengths)))
            new_tensor_tuple = tuple(x[:y]
                                     for x, y in zip(tensor_tuple, new_lens))

            for x, y in zip(new_tensor_tuple, tensor_tuple):
                x[-1] = y[-1]
            collected_tuples.append(new_tensor_tuple)
        else:
            collected_tuples.append(tensor_tuple)

    concat_tensor = list(torch.cat(x, dim=0) for x in collected_tuples)
    all_input_concat_padded = pad_sequence(
        concat_tensor, batch_first=True, padding_value=pad_idx)
    return all_input_concat_padded