Add trainer for UniTE

modelscope/metainfo.py

@@ -875,6 +875,7 @@ class NLPTrainers(object):
     document_grounded_dialog_rerank_trainer = 'document-grounded-dialog-rerank-trainer'
     document_grounded_dialog_retrieval_trainer = 'document-grounded-dialog-retrieval-trainer'
     siamese_uie_trainer = 'siamese-uie-trainer'
+    translation_evaluation_trainer = 'translation-evaluation-trainer'
 
 
 class MultiModalTrainers(object):
@@ -1089,6 +1090,8 @@ class Metrics(object):
     # metric for image-colorization task
     image_colorization_metric = 'image-colorization-metric'
     ocr_recognition_metric = 'ocr-recognition-metric'
+    # metric for translation evaluation
+    translation_evaluation_metric = 'translation-evaluation-metric'
 
 
 class Optimizers(object):

modelscope/metrics/__init__.py

@@ -31,6 +31,7 @@ if TYPE_CHECKING:
     from .loss_metric import LossMetric
     from .image_colorization_metric import ImageColorizationMetric
     from .ocr_recognition_metric import OCRRecognitionMetric
+    from .translation_evaluation_metric import TranslationEvaluationMetric
 else:
     _import_structure = {
         'audio_noise_metric': ['AudioNoiseMetric'],
@@ -62,7 +63,8 @@ else:
         'text_ranking_metric': ['TextRankingMetric'],
         'loss_metric': ['LossMetric'],
         'image_colorization_metric': ['ImageColorizationMetric'],
-        'ocr_recognition_metric': ['OCRRecognitionMetric']
+        'ocr_recognition_metric': ['OCRRecognitionMetric'],
+        'translation_evaluation_metric': ['TranslationEvaluationMetric']
     }
 
     import sys

modelscope/metrics/builder.py

@@ -42,6 +42,7 @@ class MetricKeys(object):
     NDCG = 'ndcg'
     AR = 'AR'
     Colorfulness = 'colorfulness'
+    Kendall_Tau_Correlation = 'kendall_tau_correlation'
 
 
 task_default_metrics = {
@@ -76,6 +77,7 @@ task_default_metrics = {
     Tasks.bad_image_detecting: [Metrics.accuracy],
     Tasks.ocr_recognition: [Metrics.ocr_recognition_metric],
     Tasks.efficient_diffusion_tuning: [Metrics.loss_metric],
+    Tasks.translation_evaluation: [Metrics.translation_evaluation_metric]
 }
 
 

modelscope/metrics/translation_evaluation_metric.py

@@ -0,0 +1,174 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import importlib
+from typing import Dict, List, Union
+
+from pandas import DataFrame
+
+from modelscope.metainfo import Metrics
+from modelscope.metrics.base import Metric
+from modelscope.metrics.builder import METRICS, MetricKeys
+from modelscope.models.nlp.unite.configuration import InputFormat
+from modelscope.utils.logger import get_logger
+from modelscope.utils.registry import default_group
+
+logger = get_logger()
+
+
+@METRICS.register_module(
+    group_key=default_group, module_name=Metrics.translation_evaluation_metric)
+class TranslationEvaluationMetric(Metric):
+    r"""The metric class for translation evaluation.
+
+    """
+
+    def __init__(self, gap_threshold: float = 25.0):
+        r"""Build a translation evaluation metric, following the designed
+            Kendall's tau correlation from WMT Metrics Shared Task competitions.
+
+            Args:
+                gap_threshold: The score gap denoting the available hypothesis pair.
+
+            Returns:
+                A metric for translation evaluation.
+        """
+        self.gap_threshold = gap_threshold
+
+        self.lp = list()
+        self.segment_id = list()
+        self.raw_score = list()
+        self.score = list()
+        self.input_format = list()
+
+    def clear(self) -> None:
+        r"""Clear all the stored variables.
+        """
+        self.lp.clear()
+        self.segment_id.clear()
+        self.raw_score.clear()
+        self.input_format.clear()
+
+        self.score.clear()
+
+        return
+
+    def add(self, outputs: Dict[str, List[float]],
+            inputs: Dict[str, List[Union[float, int]]]) -> None:
+        r"""Collect the related results for processing.
+
+            Args:
+                outputs: Dict containing 'scores'
+                inputs: Dict containing 'labels' and 'segment_ids'
+
+        """
+
+        self.lp += inputs['lp']
+        self.segment_id += inputs['segment_id']
+        self.raw_score += inputs['raw_score']
+        self.input_format += inputs['input_format']
+
+        self.score += outputs['score']
+
+        return
+
+    def evaluate(self) -> Dict[str, Dict[str, float]]:
+        r"""Compute the Kendall's tau correlation.
+
+            Returns:
+                A dict denoting Kendall's tau correlation.
+
+        """
+
+        data = {
+            'lp': self.lp,
+            'segment_id': self.segment_id,
+            'raw_score': self.raw_score,
+            'input_format': self.input_format,
+            'score': self.score
+        }
+        data = DataFrame(data=data)
+        correlation = dict()
+
+        for input_format in data.input_format.unique():
+            logger.info('Evaluation results for %s input format'
+                        % input_format.value)
+            input_format_data = data[data.input_format == input_format]
+
+            temp_correlation = dict()
+
+            for lp in sorted(input_format_data.lp.unique()):
+                sub_data = input_format_data[input_format_data.lp == lp]
+                temp_correlation[input_format.value + '_'
+                                 + lp] = self.compute_kendall_tau(sub_data)
+                logger.info(
+                    '\t%s: %f' %
+                    (lp,
+                     temp_correlation[input_format.value + '_' + lp] * 100))
+
+            avg_correlation = sum(
+                temp_correlation.values()) / len(temp_correlation)
+            correlation[input_format.value + '_avg'] = avg_correlation
+            logger.info('Average evaluation result for %s input format: %f' %
+                        (input_format.value, avg_correlation))
+            logger.info('')
+            correlation.update(temp_correlation)
+
+        return correlation
+
+    def merge(self, other: 'TranslationEvaluationMetric') -> None:
+        r"""Merge the predictions from other TranslationEvaluationMetric objects.
+
+            Args:
+                other: Another TranslationEvaluationMetric object.
+
+        """
+
+        self.lp += other.lp
+        self.segment_id += other.segment_ids
+        self.raw_score += other.raw_score
+        self.input_format += other.input_format
+
+        self.score += other.score
+
+        return
+
+    def compute_kendall_tau(self, csv_data: DataFrame) -> float:
+        r"""Compute kendall's tau correlation.
+
+            Args:
+                csv_data: The pandas dataframe.
+
+            Returns:
+                float: THe kendall's Tau correlation.
+
+        """
+        concor = discor = 0
+
+        for segment_id in sorted(csv_data.segment_id.unique()):
+            group_csv_data = csv_data[csv_data.segment_id == segment_id]
+
+            examples = group_csv_data.to_dict('records')
+
+            for i in range(0, len(examples)):
+                for j in range(i + 1, len(examples)):
+                    if self.raw_score[i] - self.raw_score[
+                            j] >= self.gap_threshold:
+                        if self.score[i] > self.score[j]:
+                            concor += 1
+                        elif self.score[i] < self.score[j]:
+                            discor += 1
+                    elif self.raw_score[i] - self.raw_score[
+                            j] <= -self.gap_threshold:
+                        if self.score[i] < self.score[j]:
+                            concor += 1
+                        elif self.score[i] > self.score[j]:
+                            discor += 1
+
+        if concor + discor == 0:
+            logger.warning(
+                'We don\'t have available pairs when evaluation. '
+                'Marking the kendall tau correlation as the lowest value (-1.0).'
+            )
+            return -1.0
+        else:
+            return (concor - discor) / (concor + discor)

modelscope/models/nlp/unite/__init__.py

@@ -5,12 +5,12 @@ from typing import TYPE_CHECKING
 from modelscope.utils.import_utils import LazyImportModule
 
 if TYPE_CHECKING:
-    from .configuration_unite import UniTEConfig
-    from .modeling_unite import UniTEForTranslationEvaluation
+    from .configuration import UniTEConfig
+    from .translation_evaluation import UniTEForTranslationEvaluation
 else:
     _import_structure = {
-        'configuration_unite': ['UniTEConfig'],
-        'modeling_unite': ['UniTEForTranslationEvaluation'],
+        'configuration': ['UniTEConfig'],
+        'translation_evaluation': ['UniTEForTranslationEvaluation'],
     }
 
     import sys

modelscope/models/nlp/unite/configuration.py

@@ -9,7 +9,7 @@ from modelscope.utils.config import Config
 logger = logging.get_logger()
 
 
-class EvaluationMode(Enum):
+class InputFormat(Enum):
     SRC = 'src'
     REF = 'ref'
     SRC_REF = 'src-ref'

modelscope/models/nlp/unite/translation_evaluation.py

@@ -20,6 +20,8 @@ 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.constant import Tasks
 from modelscope.utils.logger import get_logger
 
@@ -71,8 +73,16 @@ class LayerwiseAttention(Module):
         mask: torch.Tensor = None,
     ) -> torch.Tensor:
         tensors = torch.cat(list(x.unsqueeze(dim=0) for x in tensors), dim=0)
-        normed_weights = softmax(
-            self.scalar_parameters, dim=0).view(-1, 1, 1, 1)
+
+        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
@@ -97,18 +107,18 @@ class FeedForward(Module):
         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.
+            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 = []
@@ -266,8 +276,11 @@ class UniTEForTranslationEvaluation(TorchModel):
 
         return
 
-    def forward(self, input_sentences: List[torch.Tensor]):
-        input_ids = self.combine_input_sentences(input_sentences)
+    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,
@@ -276,125 +289,138 @@ class UniTEForTranslationEvaluation(TorchModel):
             return_dict=True)
         mix_states = self.layerwise_attention(outputs['hidden_states'],
                                               attention_mask)
-        pred = self.estimator(mix_states)
-        return pred.squeeze(dim=-1)
+        pred = self.estimator(mix_states).squeeze(dim=-1)
+        output = TranslationEvaluationOutput(
+            score=pred.cpu().tolist(), input_format=input_format)
 
-    def load_checkpoint(self, path: str, device: torch.device):
-        state_dict = torch.load(path, map_location=device)
-        self.load_state_dict(state_dict)
+        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)
+            self.load_state_dict(state_dict)
         logger.info('Loading checkpoint parameters from %s' % path)
         return
 
-    def combine_input_sentences(self, input_sent_groups: List[torch.Tensor]):
-        for input_sent_group in input_sent_groups[1:]:
-            input_sent_group[:, 0] = self.eos_token_id
 
-        if len(input_sent_groups) == 3:
-            cutted_sents = self.cut_long_sequences3(input_sent_groups)
-        else:
-            cutted_sents = self.cut_long_sequences2(input_sent_groups)
-        return cutted_sents
-
-    @staticmethod
-    def cut_long_sequences2(all_input_concat: List[List[torch.Tensor]],
+def combine_input_sentences(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:
-            all_lens = tuple(len(x) for x in tensor_tuple)
+                            pad_idx: int = 1,
+                            eos_idx: int = 2):
+    for group in all_input_concat[1:]:
+        group[:, 0] = eos_idx
 
-            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))
+    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)
 
-                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]]
+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)
 
-                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)
+        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:
-                collected_tuples.append(tensor_tuple)
+                lengths[lengths_sorted_idxes[0]] = maximum_length - lengths[
+                    lengths_sorted_idxes[1]]
 
-        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)
+            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)
 
-        return all_input_concat_padded
+    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
 
-    @staticmethod
-    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:
-            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))
+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)
 
-                offset = ceil((sum(lengths.values()) - maximum_length) / 3)
+        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))
 
-                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]]
+            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]] = maximum_length // 3
+                                    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))
+            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)
+            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
+    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

modelscope/outputs/nlp_outputs.py

@@ -454,3 +454,13 @@ class SentencEmbeddingModelOutput(ModelOutputBase):
     query_embeddings: Tensor = None
     doc_embeddings: Tensor = None
     loss: Tensor = None
+
+
+@dataclass
+class TranslationEvaluationOutput(ModelOutputBase):
+    """The output class for translation evaluation models.
+    """
+
+    score: Tensor = None
+    loss: Tensor = None
+    input_format: List[str] = None

modelscope/outputs/outputs.py

@@ -1447,9 +1447,9 @@ TASK_OUTPUTS = {
     # }
     Tasks.image_skychange: [OutputKeys.OUTPUT_IMG],
     # {
-    #     'scores': [0.1, 0.2, 0.3, ...]
+    #     'score': [0.1, 0.2, 0.3, ...]
     # }
-    Tasks.translation_evaluation: [OutputKeys.SCORES],
+    Tasks.translation_evaluation: [OutputKeys.SCORE],
 
     # video object segmentation result for a single video
     #   {

modelscope/pipelines/nlp/translation_evaluation_pipeline.py

@@ -9,12 +9,11 @@ import torch
 
 from modelscope.metainfo import Pipelines
 from modelscope.models.base import Model
-from modelscope.models.nlp.unite.configuration_unite import EvaluationMode
+from modelscope.models.nlp.unite.configuration import InputFormat
 from modelscope.outputs import OutputKeys
 from modelscope.pipelines.base import InputModel, Pipeline
 from modelscope.pipelines.builder import PIPELINES
-from modelscope.preprocessors import (Preprocessor,
-                                      TranslationEvaluationPreprocessor)
+from modelscope.preprocessors import Preprocessor
 from modelscope.utils.config import Config
 from modelscope.utils.constant import ModelFile, Tasks
 from modelscope.utils.logger import get_logger
@@ -31,57 +30,55 @@ class TranslationEvaluationPipeline(Pipeline):
     def __init__(self,
                  model: InputModel,
                  preprocessor: Optional[Preprocessor] = None,
-                 eval_mode: EvaluationMode = EvaluationMode.SRC_REF,
+                 input_format: InputFormat = InputFormat.SRC_REF,
                  device: str = 'gpu',
                  **kwargs):
-        r"""Build a translation pipeline with a model dir or a model id in the model hub.
+        r"""Build a translation evaluation pipeline with a model dir or a model id in the model hub.
 
         Args:
             model: A Model instance.
-            eval_mode: Evaluation mode, choosing one from `"EvaluationMode.SRC_REF"`,
-                `"EvaluationMode.SRC"`, `"EvaluationMode.REF"`. Aside from hypothesis, the
+            preprocessor: The preprocessor for this pipeline.
+            input_format: Input format, choosing one from `"InputFormat.SRC_REF"`,
+                `"InputFormat.SRC"`, `"InputFormat.REF"`. Aside from hypothesis, the
                 source/reference/source+reference can be presented during evaluation.
+            device: Used device for this pipeline.
         """
         super().__init__(model=model, preprocessor=preprocessor)
 
-        self.eval_mode = eval_mode
-        self.checking_eval_mode()
+        self.input_format = input_format
+        self.checking_input_format()
         assert isinstance(self.model, Model), \
             f'please check whether model config exists in {ModelFile.CONFIGURATION}'
 
-        self.preprocessor = TranslationEvaluationPreprocessor(
-            self.model.model_dir,
-            self.eval_mode) if preprocessor is None else preprocessor
-
         self.model.load_checkpoint(
             osp.join(self.model.model_dir, ModelFile.TORCH_MODEL_BIN_FILE),
-            self.device)
+            device=self.device,
+            plm_only=False)
         self.model.eval()
 
         return
 
-    def checking_eval_mode(self):
-        if self.eval_mode == EvaluationMode.SRC:
+    def checking_input_format(self):
+        if self.input_format == InputFormat.SRC:
             logger.info('Evaluation mode: source-only')
-        elif self.eval_mode == EvaluationMode.REF:
+        elif self.input_format == InputFormat.REF:
             logger.info('Evaluation mode: reference-only')
-        elif self.eval_mode == EvaluationMode.SRC_REF:
+        elif self.input_format == InputFormat.SRC_REF:
             logger.info('Evaluation mode: source-reference-combined')
         else:
-            raise ValueError(
-                'Evaluation mode should be one choice among'
-                '\'EvaluationMode.SRC\', \'EvaluationMode.REF\', and'
-                '\'EvaluationMode.SRC_REF\'.')
+            raise ValueError('Evaluation mode should be one choice among'
+                             '\'InputFormat.SRC\', \'InputFormat.REF\', and'
+                             '\'InputFormat.SRC_REF\'.')
 
-    def change_eval_mode(self,
-                         eval_mode: EvaluationMode = EvaluationMode.SRC_REF):
+    def change_input_format(self,
+                            input_format: InputFormat = InputFormat.SRC_REF):
         logger.info('Changing the evaluation mode.')
-        self.eval_mode = eval_mode
-        self.checking_eval_mode()
-        self.preprocessor.eval_mode = eval_mode
+        self.input_format = input_format
+        self.checking_input_format()
+        self.preprocessor.change_input_format(input_format)
         return
 
-    def __call__(self, input: Dict[str, Union[str, List[str]]], **kwargs):
+    def __call__(self, input_dict: Dict[str, Union[str, List[str]]], **kwargs):
         r"""Implementation of __call__ function.
 
         Args:
@@ -104,12 +101,12 @@ class TranslationEvaluationPipeline(Pipeline):
                 }
                 ```
         """
-        return super().__call__(input=input, **kwargs)
+        return super().__call__(input=input_dict, **kwargs)
 
-    def forward(self,
-                input_ids: List[torch.Tensor]) -> Dict[str, torch.Tensor]:
-        return self.model(input_ids)
+    def forward(
+            self, input_dict: Dict[str,
+                                   torch.Tensor]) -> Dict[str, torch.Tensor]:
+        return self.model(**input_dict)
 
     def postprocess(self, output: torch.Tensor) -> Dict[str, Any]:
-        result = {OutputKeys.SCORES: output.cpu().tolist()}
-        return result
+        return output

modelscope/preprocessors/__init__.py

@@ -41,9 +41,9 @@ if TYPE_CHECKING:
         DialogStateTrackingPreprocessor, ConversationalTextToSqlPreprocessor,
         TableQuestionAnsweringPreprocessor, NERPreprocessorViet,
         NERPreprocessorThai, WordSegmentationPreprocessorThai,
-        TranslationEvaluationPreprocessor, CanmtTranslationPreprocessor,
-        DialogueClassificationUsePreprocessor, SiameseUiePreprocessor,
-        DocumentGroundedDialogGeneratePreprocessor,
+        TranslationEvaluationTransformersPreprocessor,
+        CanmtTranslationPreprocessor, DialogueClassificationUsePreprocessor,
+        SiameseUiePreprocessor, DocumentGroundedDialogGeneratePreprocessor,
         DocumentGroundedDialogRetrievalPreprocessor,
         DocumentGroundedDialogRerankPreprocessor)
     from .video import ReadVideoData, MovieSceneSegmentationPreprocessor
@@ -96,7 +96,7 @@ else:
             'DialogStateTrackingPreprocessor',
             'ConversationalTextToSqlPreprocessor',
             'TableQuestionAnsweringPreprocessor',
-            'TranslationEvaluationPreprocessor',
+            'TranslationEvaluationTransformersPreprocessor',
             'CanmtTranslationPreprocessor',
             'DialogueClassificationUsePreprocessor', 'SiameseUiePreprocessor',
             'DialogueClassificationUsePreprocessor',

modelscope/preprocessors/nlp/__init__.py

@@ -29,7 +29,7 @@ if TYPE_CHECKING:
     from .space_T_en import ConversationalTextToSqlPreprocessor
     from .space_T_cn import TableQuestionAnsweringPreprocessor
     from .mglm_summarization_preprocessor import MGLMSummarizationPreprocessor
-    from .translation_evaluation_preprocessor import TranslationEvaluationPreprocessor
+    from .translation_evaluation_preprocessor import TranslationEvaluationTransformersPreprocessor
     from .canmt_translation import CanmtTranslationPreprocessor
     from .dialog_classification_use_preprocessor import DialogueClassificationUsePreprocessor
     from .siamese_uie_preprocessor import SiameseUiePreprocessor
@@ -90,7 +90,7 @@ else:
         'space_T_en': ['ConversationalTextToSqlPreprocessor'],
         'space_T_cn': ['TableQuestionAnsweringPreprocessor'],
         'translation_evaluation_preprocessor':
-        ['TranslationEvaluationPreprocessor'],
+        ['TranslationEvaluationTransformersPreprocessor'],
         'canmt_translation': [
             'CanmtTranslationPreprocessor',
         ],

modelscope/preprocessors/nlp/translation_evaluation_preprocessor.py

@@ -2,10 +2,13 @@
 
 from typing import Any, Dict, List, Union
 
+import torch
 from transformers import AutoTokenizer
 
 from modelscope.metainfo import Preprocessors
-from modelscope.models.nlp.unite.configuration_unite import EvaluationMode
+from modelscope.models.nlp.unite.configuration import InputFormat
+from modelscope.models.nlp.unite.translation_evaluation import \
+    combine_input_sentences
 from modelscope.preprocessors import Preprocessor
 from modelscope.preprocessors.builder import PREPROCESSORS
 from modelscope.utils.constant import Fields, ModeKeys
@@ -14,43 +17,98 @@ from .transformers_tokenizer import NLPTokenizer
 
 @PREPROCESSORS.register_module(
     Fields.nlp, module_name=Preprocessors.translation_evaluation)
-class TranslationEvaluationPreprocessor(Preprocessor):
+class TranslationEvaluationTransformersPreprocessor(Preprocessor):
     r"""The tokenizer preprocessor used for translation evaluation.
     """
 
     def __init__(self,
                  model_dir: str,
-                 eval_mode: EvaluationMode,
+                 max_len: int,
+                 pad_token_id: int,
+                 eos_token_id: int,
+                 input_format: InputFormat = InputFormat.SRC_REF,
                  mode=ModeKeys.INFERENCE,
                  *args,
                  **kwargs):
-        r"""preprocess the data via the vocab file from the `model_dir` path
+        r"""Preprocessing the data for the model in `model_dir` path
 
         Args:
             model_dir: A Model instance.
-            eval_mode: Evaluation mode, choosing one from `"EvaluationMode.SRC_REF"`,
-                `"EvaluationMode.SRC"`, `"EvaluationMode.REF"`. Aside from hypothesis, the
+            max_len: Maximum length for input sequence.
+            pad_token_id: Token id for padding token.
+            eos_token_id: Token id for the ending-of-sequence (eos) token.
+            input_format: Input format, choosing one from `"InputFormat.SRC_REF"`,
+                `"InputFormat.SRC"`, `"InputFormat.REF"`. Aside from hypothesis, the
                 source/reference/source+reference can be presented during evaluation.
+            mode: The mode for this preprocessor.
         """
         super().__init__(mode=mode)
         self.tokenizer = NLPTokenizer(
             model_dir=model_dir, use_fast=False, tokenize_kwargs=kwargs)
-        self.eval_mode = eval_mode
+        self.input_format = input_format
+
+        self.max_len = max_len
+        self.pad_token_id = pad_token_id
+        self.eos_token_id = eos_token_id
 
         return
 
-    def __call__(self, input_dict: Dict[str, Any]) -> List[List[str]]:
-        if self.eval_mode == EvaluationMode.SRC and 'src' not in input_dict.keys(
+    def change_input_format(self, input_format: InputFormat):
+        r"""Change the input format for the preprocessor.
+
+        Args:
+            input_format: Any choice in InputFormat.SRC_REF, InputFormat.SRC and InputFormat.REF.
+
+        """
+        self.input_format = input_format
+        return
+
+    def collect_input_ids(self, input_dict: Dict[str, Any]):
+        r"""Collect the input ids for the given examples.
+
+        Args:
+            input_dict: A dict containing hyp/src/ref sentences.
+
+        Returns:
+            The token ids for each example.
+
+        """
+        output_sents = [
+            self.tokenizer(
+                input_dict['hyp'], return_tensors='pt',
+                padding=True)['input_ids']
+        ]
+        if self.input_format == InputFormat.SRC or self.input_format == InputFormat.SRC_REF:
+            output_sents += [
+                self.tokenizer(
+                    input_dict['src'], return_tensors='pt',
+                    padding=True)['input_ids']
+            ]
+        if self.input_format == InputFormat.REF or self.input_format == InputFormat.SRC_REF:
+            output_sents += [
+                self.tokenizer(
+                    input_dict['ref'], return_tensors='pt',
+                    padding=True)['input_ids']
+            ]
+
+        input_ids = combine_input_sentences(output_sents, self.max_len,
+                                            self.pad_token_id,
+                                            self.eos_token_id)
+
+        return input_ids
+
+    def __call__(self, input_dict: Dict[str, Any]) -> Dict[str, Any]:
+        if self.input_format == InputFormat.SRC and 'src' not in input_dict.keys(
         ):
             raise ValueError(
                 'Source sentences are required for source-only evaluation mode.'
             )
-        if self.eval_mode == EvaluationMode.REF and 'ref' not in input_dict.keys(
+        if self.input_format == InputFormat.REF and 'ref' not in input_dict.keys(
         ):
             raise ValueError(
                 'Reference sentences are required for reference-only evaluation mode.'
             )
-        if self.eval_mode == EvaluationMode.SRC_REF and (
+        if self.input_format == InputFormat.SRC_REF and (
                 'src' not in input_dict.keys()
                 or 'ref' not in input_dict.keys()):
             raise ValueError(
@@ -59,29 +117,58 @@ class TranslationEvaluationPreprocessor(Preprocessor):
 
         if type(input_dict['hyp']) == str:
             input_dict['hyp'] = [input_dict['hyp']]
-        if (self.eval_mode == EvaluationMode.SRC or self.eval_mode
-                == EvaluationMode.SRC_REF) and type(input_dict['src']) == str:
+        if (self.input_format == InputFormat.SRC or self.input_format
+                == InputFormat.SRC_REF) and type(input_dict['src']) == str:
             input_dict['src'] = [input_dict['src']]
-        if (self.eval_mode == EvaluationMode.REF or self.eval_mode
-                == EvaluationMode.SRC_REF) and type(input_dict['ref']) == str:
+        if (self.input_format == InputFormat.REF or self.input_format
+                == InputFormat.SRC_REF) and type(input_dict['ref']) == str:
             input_dict['ref'] = [input_dict['ref']]
 
-        output_sents = [
-            self.tokenizer(
-                input_dict['hyp'], return_tensors='pt',
-                padding=True)['input_ids']
-        ]
-        if self.eval_mode == EvaluationMode.SRC or self.eval_mode == EvaluationMode.SRC_REF:
-            output_sents += [
-                self.tokenizer(
-                    input_dict['src'], return_tensors='pt',
-                    padding=True)['input_ids']
-            ]
-        if self.eval_mode == EvaluationMode.REF or self.eval_mode == EvaluationMode.SRC_REF:
-            output_sents += [
-                self.tokenizer(
-                    input_dict['ref'], return_tensors='pt',
-                    padding=True)['input_ids']
-            ]
+        if (self.input_format == InputFormat.SRC
+                or self.input_format == InputFormat.SRC_REF) and (len(
+                    input_dict['hyp']) != len(input_dict['src'])):
+            raise ValueError(
+                'The number of given hyp sentences (%d) is not equal to that of src (%d).'
+                % (len(input_dict['hyp']), len(input_dict['src'])))
+        if (self.input_format == InputFormat.REF
+                or self.input_format == InputFormat.SRC_REF) and (len(
+                    input_dict['hyp']) != len(input_dict['ref'])):
+            raise ValueError(
+                'The number of given hyp sentences (%d) is not equal to that of ref (%d).'
+                % (len(input_dict['hyp']), len(input_dict['ref'])))
 
-        return output_sents
+        output_dict = {'input_ids': self.collect_input_ids(input_dict)}
+
+        if self.mode == ModeKeys.TRAIN or self.mode == ModeKeys.EVAL:
+            if 'score' not in input_dict.keys():
+                raise KeyError(
+                    'During training or evaluating, \'score\' should be provided.'
+                )
+            if (isinstance(input_dict['score'], List) and len(input_dict['score']) != len(output_dict['input_ids'])) \
+                    or (isinstance(input_dict['score'], float) and len(output['input_ids']) != 1):
+                raise ValueError(
+                    'The number of score is not equal to that of the given examples. '
+                    'Required %d, given %d.' %
+                    (len(output['input_ids']), len(input_dict['score'])))
+
+            output_dict['score'] = [input_dict['score']] if isinstance(
+                input_dict['score'], float) else input_dict['score']
+
+        if self.mode == ModeKeys.EVAL:
+            if 'lp' not in input_dict.keys():
+                raise ValueError(
+                    'Language pair should be provided for evaluation.')
+
+            if 'segment_id' not in input_dict.keys():
+                raise ValueError(
+                    'Segment id should be provided for evaluation.')
+
+            if 'raw_score' not in input_dict.keys():
+                raise ValueError(
+                    'Raw scores should be provided for evaluation.')
+
+            output_dict['lp'] = input_dict['lp']
+            output_dict['segment_id'] = input_dict['segment_id']
+            output_dict['raw_score'] = input_dict['raw_score']
+
+        return output_dict

modelscope/trainers/nlp/__init__.py

@@ -10,6 +10,7 @@ if TYPE_CHECKING:
     from .text_generation_trainer import TextGenerationTrainer
     from .sentence_embedding_trainer import SentenceEmbeddingTrainer
     from .siamese_uie_trainer import SiameseUIETrainer
+    from .translation_evaluation_trainer import TranslationEvaluationTrainer
 else:
     _import_structure = {
         'sequence_classification_trainer': ['SequenceClassificationTrainer'],
@@ -17,7 +18,8 @@ else:
         'text_ranking_trainer': ['TextRankingTrainer'],
         'text_generation_trainer': ['TextGenerationTrainer'],
         'sentence_emebedding_trainer': ['SentenceEmbeddingTrainer'],
-        'siamese_uie_trainer': ['SiameseUIETrainer']
+        'siamese_uie_trainer': ['SiameseUIETrainer'],
+        'translation_evaluation_trainer': ['TranslationEvaluationTrainer']
     }
 
     import sys

modelscope/trainers/nlp/translation_evaluation_trainer.py

@@ -0,0 +1,396 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+"""PyTorch trainer for UniTE model."""
+
+import os.path as osp
+import random
+from math import ceil
+from os import mkdir
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import torch
+from pandas import DataFrame
+from torch.nn.functional import pad
+from torch.nn.utils import clip_grad_norm_
+from torch.optim import AdamW, Optimizer
+from torch.utils.data import (BatchSampler, DataLoader, Dataset, Sampler,
+                              SequentialSampler, SubsetRandomSampler)
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+from transformers import AutoTokenizer
+
+from modelscope.metainfo import Metrics, Trainers
+from modelscope.metrics import Metric
+from modelscope.metrics.builder import MetricKeys, build_metric
+from modelscope.models.base import TorchModel
+from modelscope.models.nlp.unite.configuration import InputFormat
+from modelscope.models.nlp.unite.translation_evaluation import (
+    UniTEForTranslationEvaluation, combine_input_sentences)
+from modelscope.msdatasets import MsDataset
+from modelscope.preprocessors import Preprocessor
+from modelscope.trainers.builder import TRAINERS
+from modelscope.trainers.hooks import Hook
+from modelscope.trainers.trainer import EpochBasedTrainer
+from modelscope.utils.config import ConfigDict
+from modelscope.utils.constant import (ConfigKeys, Fields, ModeKeys, ModelFile,
+                                       TrainerStages)
+from modelscope.utils.device import create_device
+from modelscope.utils.logger import get_logger
+
+logger = get_logger()
+
+
+class TranslationEvaluationTrainingSampler(Sampler):
+
+    def __init__(self, num_of_samples: int,
+                 batch_size_for_each_input_format: int):
+        r"""Build a sampler for model training with translation evaluation trainer.
+        The trainer should derive samples for each subset of the entire dataset.
+
+        Args:
+            num_of_samples: The number of samples in total.
+            batch_size_for_each_input_format: During training, the batch size for each input format
+
+        Returns:
+            A data sampler for translation evaluation model training.
+
+        """
+
+        self.num_of_samples = num_of_samples
+        self.batch_size_for_each_input_format = batch_size_for_each_input_format
+
+        self.num_of_samples_for_each_input_format = self.num_of_samples // 3
+        num_of_samples_to_use = self.num_of_samples_for_each_input_format * 3
+
+        logger.info(
+            '%d samples are given for training. '
+            'Using %d samples for each input format. '
+            'Leaving the last %d samples unused.' %
+            (self.num_of_samples, self.num_of_samples_for_each_input_format,
+             self.num_of_samples - num_of_samples_to_use))
+        self.num_of_samples = num_of_samples_to_use
+
+        random_permutations = torch.randperm(
+            self.num_of_samples).cpu().tolist()
+
+        self.subset_iterators = dict()
+        self.subset_samplers = dict()
+        self.indices_for_each_input_format = dict()
+        for input_format_index, input_format in \
+                enumerate((InputFormat.SRC_REF, InputFormat.SRC, InputFormat.REF)):
+            start_idx = input_format_index * self.num_of_samples_for_each_input_format
+            end_idx = start_idx + self.num_of_samples_for_each_input_format
+            self.indices_for_each_input_format[
+                input_format] = random_permutations[start_idx:end_idx]
+            self.subset_samplers[input_format] = \
+                BatchSampler(SubsetRandomSampler(self.indices_for_each_input_format[input_format]),
+                             batch_size=self.batch_size_for_each_input_format,
+                             drop_last=True)
+            self.subset_iterators[input_format] = iter(
+                self.subset_samplers[input_format])
+
+        self.num_of_sampled_batches = 0
+
+        if self.__len__() == 0:
+            raise ValueError(
+                'The dataset doesn\'t contain enough examples to form a single batch.',
+                'Please reduce the batch_size or use more examples for training.'
+            )
+
+        return
+
+    def __iter__(self):
+        while True:
+            try:
+                if self.num_of_sampled_batches == self.__len__():
+                    for input_format in (InputFormat.SRC_REF, InputFormat.SRC,
+                                         InputFormat.REF):
+                        while True:
+                            try:
+                                next(self.subset_iterators[input_format])
+                            except StopIteration:
+                                self.subset_iterators[input_format] = \
+                                    iter(self.subset_samplers[input_format])
+                                break
+
+                    self.num_of_sampled_batches = 0
+
+                output = list()
+                for input_format_idx, input_format in \
+                        enumerate((InputFormat.SRC_REF, InputFormat.SRC, InputFormat.REF)):
+                    output += next(self.subset_iterators[input_format])
+
+                self.num_of_sampled_batches += 1
+
+                yield output
+            except StopIteration:
+                break
+
+    def __len__(self) -> int:
+        return self.num_of_samples_for_each_input_format // self.batch_size_for_each_input_format
+
+
+def convert_csv_dict_to_input(
+        batch: List[Dict[str, Any]],
+        preprocessor: Preprocessor) -> Tuple[List[torch.Tensor]]:
+
+    input_dict = dict()
+
+    for key in batch[0].keys():
+        input_dict[key] = list(x[key] for x in batch)
+
+    input_dict = preprocessor(input_dict)
+
+    return input_dict
+
+
+def data_collate_fn(batch: List[Dict[str, Any]], batch_size: int,
+                    preprocessor: Preprocessor) -> List[Dict[str, Any]]:
+
+    output_dict = dict()
+    output_dict['input_format'] = list()
+
+    if preprocessor.mode == ModeKeys.TRAIN:
+        for input_format_index, input_format in \
+                enumerate((InputFormat.SRC_REF, InputFormat.SRC, InputFormat.REF)):
+            start_idx = input_format_index * batch_size
+            end_idx = start_idx + batch_size
+            batch_to_process = batch[start_idx:end_idx]
+            output_dict['input_format'] += [input_format] * batch_size
+            preprocessor.change_input_format(input_format)
+            batch_to_process = convert_csv_dict_to_input(
+                batch_to_process, preprocessor)
+
+            for key, value in batch_to_process.items():
+                if key not in output_dict.keys():
+                    output_dict[key] = list()
+                output_dict[key].append(value)
+    elif preprocessor.mode == ModeKeys.EVAL:
+        output_dict['input_format'] += [preprocessor.input_format] * len(batch)
+        batch = convert_csv_dict_to_input(batch, preprocessor)
+
+        for key, value in batch.items():
+            if key not in output_dict.keys():
+                output_dict[key] = list()
+            output_dict[key].append(value)
+    else:
+        raise ValueError(
+            'During training, %s mode is not allowed for preprocessor.'
+            % preprocessor.mode)
+
+    input_max_lengths = max(x.size(-1) for x in output_dict['input_ids'])
+    output_dict['input_ids'] = list(
+        pad(x,
+            pad=(0, input_max_lengths - x.size(-1)),
+            value=preprocessor.pad_token_id) for x in output_dict['input_ids'])
+
+    output_dict['input_ids'] = torch.cat(output_dict['input_ids'], dim=0)
+    output_dict['score'] = torch.Tensor(output_dict['score']).view(-1)
+
+    if preprocessor.mode == ModeKeys.EVAL:
+        output_dict['lp'] = sum(output_dict['lp'], list())
+        output_dict['raw_score'] = sum(output_dict['raw_score'], list())
+        output_dict['segment_id'] = sum(output_dict['segment_id'], list())
+
+    return output_dict
+
+
+@TRAINERS.register_module(module_name=Trainers.translation_evaluation_trainer)
+class TranslationEvaluationTrainer(EpochBasedTrainer):
+
+    def __init__(self,
+                 model: Optional[Union[TorchModel, torch.nn.Module,
+                                       str]] = None,
+                 cfg_file: Optional[str] = None,
+                 device: str = 'gpu',
+                 *args,
+                 **kwargs):
+        r"""Build a translation evaluation trainer with a model dir or a model id in the model hub.
+
+        Args:
+            model: A Model instance.
+            cfg_file: The path for the configuration file (configuration.json).
+            device: Used device for this trainer.
+
+        """
+
+        def data_collator_for_train(x):
+            return data_collate_fn(
+                x,
+                batch_size=self.cfg.train.batch_size,
+                preprocessor=self.train_preprocessor)
+
+        def data_collator_for_eval(x):
+            return data_collate_fn(
+                x,
+                batch_size=self.cfg.evaluation.batch_size,
+                preprocessor=self.eval_preprocessor)
+
+        data_collator = {
+            ConfigKeys.train: data_collator_for_train,
+            ConfigKeys.val: data_collator_for_eval
+        }
+
+        super().__init__(
+            model,
+            cfg_file=cfg_file,
+            data_collator=data_collator,
+            *args,
+            **kwargs)
+
+        self.train_dataloader = None
+        self.eval_dataloader = None
+
+        return
+
+    def build_optimizer(self, cfg: ConfigDict) -> Optimizer:
+        r"""Sets the optimizers to be used during training."""
+        if self.cfg.train.optimizer.type != 'AdamW':
+            return super().build_optimizer(cfg)
+
+        # Freezing embedding layers for more efficient training.
+        for param in self.model.encoder.embeddings.parameters():
+            param.requires_grad = False
+
+        logger.info('Building AdamW optimizer ...')
+        learning_rates_and_parameters = list({
+            'params':
+            self.model.encoder.encoder.layer[i].parameters(),
+            'lr':
+            self.cfg.train.optimizer.plm_lr
+            * self.cfg.train.optimizer.plm_lr_layerwise_decay**i,
+        } for i in range(0, self.cfg.model.num_hidden_layers))
+
+        learning_rates_and_parameters.append({
+            'params':
+            self.model.encoder.embeddings.parameters(),
+            'lr':
+            self.cfg.train.optimizer.plm_lr,
+        })
+
+        learning_rates_and_parameters.append({
+            'params':
+            self.model.estimator.parameters(),
+            'lr':
+            self.cfg.train.optimizer.mlp_lr
+        })
+
+        learning_rates_and_parameters.append({
+            'params':
+            self.model.layerwise_attention.parameters(),
+            'lr':
+            self.cfg.train.optimizer.mlp_lr,
+        })
+
+        optimizer = AdamW(
+            learning_rates_and_parameters,
+            lr=self.cfg.train.optimizer.plm_lr,
+            betas=self.cfg.train.optimizer.betas,
+            eps=self.cfg.train.optimizer.eps,
+            weight_decay=self.cfg.train.optimizer.weight_decay,
+        )
+
+        return optimizer
+
+    def get_train_dataloader(self) -> DataLoader:
+        logger.info('Building dataloader for training ...')
+
+        if self.train_dataset is None:
+            logger.info('Reading train csv file from %s ...'
+                        % self.cfg.dataset.train.name)
+            self.train_dataset = MsDataset.load(
+                osp.join(self.model_dir, self.cfg.dataset.train.name),
+                split=self.cfg.dataset.train.split)
+
+        train_dataloader = DataLoader(
+            self.train_dataset,
+            batch_sampler=TranslationEvaluationTrainingSampler(
+                len(self.train_dataset),
+                batch_size_for_each_input_format=self.cfg.train.batch_size),
+            num_workers=4,
+            collate_fn=self.train_data_collator,
+            generator=None)
+
+        logger.info('Reading done, %d items in total'
+                    % len(self.train_dataset))
+
+        return train_dataloader
+
+    def get_eval_data_loader(self) -> DataLoader:
+        logger.info('Building dataloader for evaluating ...')
+
+        if self.eval_dataset is None:
+            logger.info('Reading eval csv file from %s ...'
+                        % self.cfg.dataset.valid.name)
+
+            self.eval_dataset = MsDataset.load(
+                osp.join(self.model_dir, self.cfg.dataset.valid.name),
+                split=self.cfg.dataset.valid.split)
+
+        eval_dataloader = DataLoader(
+            self.eval_dataset,
+            batch_sampler=BatchSampler(
+                SequentialSampler(range(0, len(self.eval_dataset))),
+                batch_size=self.cfg.evaluation.batch_size,
+                drop_last=False),
+            num_workers=4,
+            collate_fn=self.eval_data_collator,
+            generator=None)
+
+        logger.info('Reading done, %d items in total' % len(self.eval_dataset))
+
+        return eval_dataloader
+
+    def evaluation_loop(self, data_loader, metric_classes):
+        """ Evaluation loop used by `TranslationEvaluationTrainer.evaluate()`.
+
+        The evaluation process of UniTE model should be arranged with three loops,
+        corresponding to the input formats of `InputFormat.SRC_REF`, `InputFormat.REF`,
+        and `InputFormat.SRC`.
+
+        Here we directly copy the codes of `EpochBasedTrainer.evaluation_loop`, and change
+        the input format during each evaluation subloop.
+        """
+        vis_closure = None
+        if hasattr(self.cfg.evaluation, 'visualization'):
+            vis_cfg = self.cfg.evaluation.visualization
+            vis_closure = partial(
+                self.visualization, dataset=self.eval_dataset, **vis_cfg)
+
+        self.invoke_hook(TrainerStages.before_val)
+        metric_values = dict()
+
+        for input_format in (InputFormat.SRC_REF, InputFormat.SRC,
+                             InputFormat.REF):
+            self.eval_preprocessor.change_input_format(input_format)
+
+            if self._dist:
+                from modelscope.trainers.utils.inference import multi_gpu_test
+                # list of batched result and data samples
+                metric_values.update(
+                    multi_gpu_test(
+                        self,
+                        data_loader,
+                        device=self.device,
+                        metric_classes=metric_classes,
+                        vis_closure=vis_closure,
+                        tmpdir=self.cfg.evaluation.get('cache_dir', None),
+                        gpu_collect=self.cfg.evaluation.get(
+                            'gpu_collect', False),
+                        data_loader_iters_per_gpu=self._eval_iters_per_epoch))
+            else:
+                from modelscope.trainers.utils.inference import single_gpu_test
+                metric_values.update(
+                    single_gpu_test(
+                        self,
+                        data_loader,
+                        device=self.device,
+                        metric_classes=metric_classes,
+                        vis_closure=vis_closure,
+                        data_loader_iters=self._eval_iters_per_epoch))
+
+            for m in metric_classes:
+                if hasattr(m, 'clear') and callable(m.clear):
+                    m.clear()
+
+        self.invoke_hook(TrainerStages.after_val)
+        return metric_values

tests/metrics/test_translation_evaluation_metrics.py

@@ -0,0 +1,30 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import unittest
+
+from modelscope.metrics.translation_evaluation_metric import \
+    TranslationEvaluationMetric
+from modelscope.models.nlp.unite.configuration import InputFormat
+from modelscope.utils.test_utils import test_level
+
+
+class TestTranslationEvaluationMetrics(unittest.TestCase):
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_value(self):
+        metric = TranslationEvaluationMetric(gap_threshold=25.0)
+
+        outputs = {'score': [0.25, 0.22, 0.30, 0.78, 1.11, 0.95, 1.00, 0.86]}
+        inputs = {
+            'lp': ['zh-en'] * 8,
+            'segment_id': [0, 0, 0, 1, 1, 2, 2, 2],
+            'raw_score': [94.0, 60.0, 25.0, 59.5, 90.0, 100.0, 80.0, 60.0],
+            'input_format': [InputFormat.SRC_REF] * 8,
+        }
+        metric.add(outputs, inputs)
+        result = metric.evaluate()
+        print(result)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/pipelines/test_translation_evaluation.py

@@ -2,7 +2,7 @@
 
 import unittest
 
-from modelscope.models.nlp.unite.configuration_unite import EvaluationMode
+from modelscope.models.nlp.unite.configuration import InputFormat
 from modelscope.pipelines import pipeline
 from modelscope.utils.constant import Tasks
 from modelscope.utils.demo_utils import DemoCompatibilityCheck
@@ -18,7 +18,7 @@ class TranslationEvaluationTest(unittest.TestCase, DemoCompatibilityCheck):
 
     @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
     def test_run_with_model_name_for_unite_large(self):
-        input = {
+        input_dict = {
             'hyp': [
                 'This is a sentence.',
                 'This is another sentence.',
@@ -34,27 +34,27 @@ class TranslationEvaluationTest(unittest.TestCase, DemoCompatibilityCheck):
         }
 
         pipeline_ins = pipeline(self.task, model=self.model_id_large)
-        print(pipeline_ins(input=input))
+        print(pipeline_ins(input_dict)['score'])
 
-        pipeline_ins.change_eval_mode(eval_mode=EvaluationMode.SRC)
-        print(pipeline_ins(input=input))
+        pipeline_ins.change_input_format(input_format=InputFormat.SRC)
+        print(pipeline_ins(input_dict)['score'])
 
-        pipeline_ins.change_eval_mode(eval_mode=EvaluationMode.REF)
-        print(pipeline_ins(input=input))
+        pipeline_ins.change_input_format(input_format=InputFormat.REF)
+        print(pipeline_ins(input_dict)['score'])
 
         pipeline_ins = pipeline(
             self.task, model=self.model_id_large, device='cpu')
-        print(pipeline_ins(input=input))
+        print(pipeline_ins(input_dict)['score'])
 
-        pipeline_ins.change_eval_mode(eval_mode=EvaluationMode.SRC)
-        print(pipeline_ins(input=input))
+        pipeline_ins.change_input_format(input_format=InputFormat.SRC)
+        print(pipeline_ins(input_dict)['score'])
 
-        pipeline_ins.change_eval_mode(eval_mode=EvaluationMode.REF)
-        print(pipeline_ins(input=input))
+        pipeline_ins.change_input_format(input_format=InputFormat.REF)
+        print(pipeline_ins(input_dict)['score'])
 
     @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
     def test_run_with_model_name_for_unite_base(self):
-        input = {
+        input_dict = {
             'hyp': [
                 'This is a sentence.',
                 'This is another sentence.',
@@ -70,23 +70,23 @@ class TranslationEvaluationTest(unittest.TestCase, DemoCompatibilityCheck):
         }
 
         pipeline_ins = pipeline(self.task, model=self.model_id_base)
-        print(pipeline_ins(input=input))
+        print(pipeline_ins(input_dict)['score'])
 
-        pipeline_ins.change_eval_mode(eval_mode=EvaluationMode.SRC)
-        print(pipeline_ins(input=input))
+        pipeline_ins.change_input_format(input_format=InputFormat.SRC)
+        print(pipeline_ins(input_dict)['score'])
 
-        pipeline_ins.change_eval_mode(eval_mode=EvaluationMode.REF)
-        print(pipeline_ins(input=input))
+        pipeline_ins.change_input_format(input_format=InputFormat.REF)
+        print(pipeline_ins(input_dict)['score'])
 
         pipeline_ins = pipeline(
             self.task, model=self.model_id_base, device='cpu')
-        print(pipeline_ins(input=input))
+        print(pipeline_ins(input_dict)['score'])
 
-        pipeline_ins.change_eval_mode(eval_mode=EvaluationMode.SRC)
-        print(pipeline_ins(input=input))
+        pipeline_ins.change_input_format(input_format=InputFormat.SRC)
+        print(pipeline_ins(input_dict)['score'])
 
-        pipeline_ins.change_eval_mode(eval_mode=EvaluationMode.REF)
-        print(pipeline_ins(input=input))
+        pipeline_ins.change_input_format(input_format=InputFormat.REF)
+        print(pipeline_ins(input_dict)['score'])
 
 
 if __name__ == '__main__':

tests/run_config.yaml

@@ -21,6 +21,7 @@ isolated:  # test cases that may require excessive anmount of GPU memory or run
   - test_image_instance_segmentation_trainer.py
   - test_image_portrait_enhancement_trainer.py
   - test_translation_trainer.py
+  - test_translation_evaluation_trainer.py
   - test_unifold.py
   - test_automatic_post_editing.py
   - test_mplug_tasks.py
@@ -77,6 +78,7 @@ envs:
       - test_text_to_speech.py
       - test_csanmt_translation.py
       - test_translation_trainer.py
+      - test_translation_evaluation_trainer.py
       - test_ocr_detection.py
       - test_automatic_speech_recognition.py
       - test_image_matting.py

tests/trainers/model_trainer_map.py

@@ -124,6 +124,12 @@ model_trainer_map = {
     'damo/nlp_csanmt_translation_en2es': [
         'tests/trainers/test_translation_trainer.py'
     ],
+    'damo/nlp_unite_mup_translation_evaluation_multilingual_base': [
+        'tests/trainers/test_translation_evaluation_trainer.py'
+    ],
+    'damo/nlp_unite_mup_translation_evaluation_multilingual_large': [
+        'tests/trainers/test_translation_evaluation_trainer.py'
+    ],
     'damo/cv_googlenet_pgl-video-summarization': [
         'tests/trainers/test_video_summarization_trainer.py'
     ],

tests/trainers/test_translation_evaluation_trainer.py

@@ -0,0 +1,30 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import unittest
+
+from modelscope.metainfo import Trainers
+from modelscope.trainers import build_trainer
+from modelscope.utils.test_utils import test_level
+
+
+class TranslationEvaluationTest(unittest.TestCase):
+
+    def setUp(self) -> None:
+        self.name = Trainers.translation_evaluation_trainer
+        self.model_id_large = 'damo/nlp_unite_mup_translation_evaluation_multilingual_large'
+        self.model_id_base = 'damo/nlp_unite_mup_translation_evaluation_multilingual_base'
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_run_with_unite_mup_large(self) -> None:
+        default_args = {'model': self.model_id_large}
+        trainer = build_trainer(name=self.name, default_args=default_args)
+        trainer.train()
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_run_with_unite_mup_base(self) -> None:
+        default_args = {'model': self.model_id_base}
+        trainer = build_trainer(name=self.name, default_args=default_args)
+        trainer.train()
+
+
+if __name__ == '__main__':
+    unittest.main()