add pipelines for en2zh-imt and zh2en-imt

*[ 交互式机器翻译-英中-通用领域-large](https://modelscope.cn/models/damo/nlp_imt_translation_en2zh/summary)
*[ 交互式机器翻译-中英-通用领域-large](https://modelscope.cn/models/damo/nlp_imt_translation_zh2en/summary)

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11593182

modelscope/metainfo.py

@@ -393,6 +393,7 @@ class Pipelines(object):
     fill_mask = 'fill-mask'
     fill_mask_ponet = 'fill-mask-ponet'
     csanmt_translation = 'csanmt-translation'
+    interactive_translation = 'interactive-translation'
     nli = 'nli'
     dialog_intent_prediction = 'dialog-intent-prediction'
     dialog_modeling = 'dialog-modeling'

modelscope/models/nlp/csanmt/translation.py

@@ -25,14 +25,20 @@ class CsanmtForTranslation(Model):
         """
         super().__init__(model_dir, *args, **kwargs)
         self.params = kwargs
+        print(self.params)
 
     def __call__(self,
                  input: Dict[str, Tensor],
-                 label: Dict[str, Tensor] = None) -> Dict[str, Tensor]:
+                 label: Dict[str, Tensor] = None,
+                 prefix: Dict[str, Tensor] = None,
+                 prefix_hit: Dict[bool, Tensor] = None) -> Dict[str, Tensor]:
         """return the result by the model
 
         Args:
-            input: the preprocessed data
+            input: the preprocessed input source sequence
+            label: the ground truth target data for model training
+            prefix: the preprocessed input target prefix sequence for interactive translation
+            prefix_hit: the preprocessed target prefix subword vector for interactive translation
 
         Returns:
             output_seqs: output sequence of target ids
@@ -40,7 +46,11 @@ class CsanmtForTranslation(Model):
         if label is None:
             with tf.compat.v1.variable_scope('NmtModel'):
                 output_seqs, output_scores = self.beam_search(
-                    input, self.params)
+                    {
+                        'input_wids': input,
+                        'prefix_wids': prefix,
+                        'prefix_hit': prefix_hit
+                    }, self.params)
             return {
                 'output_seqs': output_seqs,
                 'output_scores': output_scores,
@@ -441,7 +451,8 @@ class CsanmtForTranslation(Model):
                        trg_seq,
                        states_key,
                        states_val,
-                       params={}):
+                       params={},
+                       is_prefix=False):
         trg_vocab_size = params['trg_vocab_size']
         hidden_size = params['hidden_size']
 
@@ -468,9 +479,10 @@ class CsanmtForTranslation(Model):
             tensor=decoder_input, paddings=[[0, 0], [1, 0], [0, 0]])[:, :-1, :]
         if params['position_info_type'] == 'absolute':
             decoder_input = add_timing_signal(decoder_input)
-
-        decoder_input = decoder_input[:, -1:, :]
-        decoder_self_attention_bias = decoder_self_attention_bias[:, :, -1:, :]
+        if not is_prefix:
+            decoder_input = decoder_input[:, -1:, :]
+            decoder_self_attention_bias = decoder_self_attention_bias[:, :,
+                                                                      -1:, :]
         decoder_output, attention_weights = transformer_decoder(
             decoder_input,
             encoder_output,
@@ -480,8 +492,12 @@ class CsanmtForTranslation(Model):
             states_val=states_val,
             embedding_augmentation=feature_output,
             params=params)
-        decoder_output_last = decoder_output[:, -1, :]
-        attention_weights_last = attention_weights[:, -1, :]
+        if not is_prefix:
+            decoder_output_last = decoder_output[:, -1, :]
+            attention_weights_last = attention_weights[:, -1, :]
+        else:
+            decoder_output_last = decoder_output
+            attention_weights_last = attention_weights
 
         if params['shared_embedding_and_softmax_weights']:
             embedding_scope = \
@@ -502,27 +518,28 @@ class CsanmtForTranslation(Model):
         num_decoder_layers = params['num_decoder_layers']
         lp_rate = params['lp_rate']
         max_decoded_trg_len = params['max_decoded_trg_len']
-        batch_size = tf.shape(input=features)[0]
+        src_input = features['input_wids']
+        if 'prefix_wids' in features:
+            prefix = features['prefix_wids']
+            prefix_hit = features['prefix_hit']
+        else:
+            prefix = None
+            prefix_hit = None
+        batch_size = tf.shape(src_input)[0]
 
-        features = tile_to_beam_size(features, beam_size)
-        features = merge_first_two_dims(features)
+        src_input = tile_to_beam_size(src_input, beam_size)
+        src_input = merge_first_two_dims(src_input)
+        if prefix is not None:
+            prefix = tf.cast(tile_to_beam_size(prefix, beam_size), tf.int32)
+            prefix_hit = tile_to_beam_size(prefix_hit, beam_size)
 
         encoder_output, encoder_self_attention_bias = self.encoding_graph(
-            features, params)
+            src_input, params)
         source_name = 'source'
         if params['shared_source_target_embedding']:
             source_name = None
         feature_output = self.semantic_encoding_graph(
-            features, params, name=source_name)
-
-        init_seqs = tf.fill([batch_size, beam_size, 1], 0)
-        init_log_probs = \
-            tf.constant([[0.] + [tf.float32.min] * (beam_size - 1)])
-        init_log_probs = tf.tile(init_log_probs, [batch_size, 1])
-        init_scores = tf.zeros_like(init_log_probs)
-        fin_seqs = tf.cast(tf.fill([batch_size, beam_size, 1], 0), tf.int32)
-        fin_scores = tf.fill([batch_size, beam_size], tf.float32.min)
-        fin_flags = tf.cast(tf.fill([batch_size, beam_size], 0), tf.bool)
+            src_input, params, name=source_name)
 
         states_key = [
             tf.fill([batch_size, 0, hidden_size], 0.0)
@@ -545,6 +562,66 @@ class CsanmtForTranslation(Model):
             tile_to_beam_size(states_val[layer], beam_size)
             for layer in range(num_decoder_layers)
         ]
+        fixed_length = 1
+        if prefix is not None:
+            init_seqs = tf.concat(
+                [prefix, tf.fill([batch_size, beam_size, 1], 0)], axis=2)
+            fixed_length = tf.shape(init_seqs)[-1]
+            flat_seqs = merge_first_two_dims(init_seqs)
+            flat_states_key = [
+                merge_first_two_dims(states_key[layer])
+                for layer in range(num_decoder_layers)
+            ]
+            flat_states_val = [
+                merge_first_two_dims(states_val[layer])
+                for layer in range(num_decoder_layers)
+            ]
+
+            step_log_probs, step_attn_weights, step_states_key, step_states_val = self.inference_func(
+                encoder_output,
+                feature_output,
+                encoder_self_attention_bias,
+                flat_seqs,
+                flat_states_key,
+                flat_states_val,
+                params=params,
+                is_prefix=True)
+
+            states_key = [
+                split_first_two_dims(step_states_key[layer], batch_size,
+                                     beam_size)
+                for layer in range(num_decoder_layers)
+            ]
+            states_val = [
+                split_first_two_dims(step_states_val[layer], batch_size,
+                                     beam_size)
+                for layer in range(num_decoder_layers)
+            ]
+
+            prefix_hit = merge_first_two_dims(prefix_hit)
+            log_probs = tf.where(
+                prefix_hit, step_log_probs[:, -1, :],
+                tf.ones_like(step_log_probs[:, -1, :]) * tf.float32.min)
+
+            init_seqs = tf.concat([
+                flat_seqs[:, :-1],
+                tf.expand_dims(
+                    tf.cast(tf.argmax(log_probs, -1), tf.int32), -1)
+            ], -1)
+
+            init_seqs = split_first_two_dims(init_seqs, batch_size, beam_size)
+            init_seqs = tf.concat(
+                [init_seqs, tf.fill([batch_size, beam_size, 1], 0)], axis=2)
+        else:
+            init_seqs = tf.fill([batch_size, beam_size, 1], 0)
+
+        init_log_probs = \
+            tf.constant([[0.] + [tf.float32.min] * (beam_size - 1)])
+        init_log_probs = tf.tile(init_log_probs, [batch_size, 1])
+        init_scores = tf.zeros_like(init_log_probs)
+        fin_seqs = init_seqs
+        fin_scores = tf.fill([batch_size, beam_size], tf.float32.min)
+        fin_flags = tf.cast(tf.fill([batch_size, beam_size], 0), tf.bool)
 
         state = BeamSearchState(
             inputs=(init_seqs, init_log_probs, init_scores),
@@ -573,7 +650,8 @@ class CsanmtForTranslation(Model):
                 flat_seqs,
                 flat_states_key,
                 flat_states_val,
-                params=params)
+                params=params,
+                is_prefix=False)
 
             step_log_probs = split_first_two_dims(step_log_probs, batch_size,
                                                   beam_size)
@@ -737,7 +815,7 @@ class CsanmtForTranslation(Model):
             tf.reduce_any(input_tensor=final_flags, axis=1), final_scores,
             alive_scores)
 
-        final_seqs = final_seqs[:, :, :-1]
+        final_seqs = final_seqs[:, :, fixed_length - 1:-1]
         return final_seqs, final_scores
 
 
@@ -936,7 +1014,7 @@ def transformer_encoder(encoder_input,
     layer_postproc = params['layer_postproc']
     x = encoder_input
     mask = tf.expand_dims(mask, 2)
-    with tf.compat.v1.variable_scope(name):
+    with tf.compat.v1.variable_scope(name, reuse=tf.compat.v1.AUTO_REUSE):
         for layer in range(num_encoder_layers):
             with tf.compat.v1.variable_scope('layer_%d' % layer):
                 max_relative_dis = params['max_relative_dis'] \
@@ -1032,7 +1110,7 @@ def transformer_decoder(decoder_input,
     layer_preproc = params['layer_preproc']
     layer_postproc = params['layer_postproc']
     x = decoder_input
-    with tf.compat.v1.variable_scope(name):
+    with tf.compat.v1.variable_scope(name, reuse=tf.compat.v1.AUTO_REUSE):
         for layer in range(num_decoder_layers):
             with tf.compat.v1.variable_scope('layer_%d' % layer):
                 max_relative_dis = params['max_relative_dis'] \

modelscope/pipelines/nlp/__init__.py

@@ -17,6 +17,7 @@ if TYPE_CHECKING:
     from .feature_extraction_pipeline import FeatureExtractionPipeline
     from .fill_mask_pipeline import FillMaskPipeline
     from .information_extraction_pipeline import InformationExtractionPipeline
+    from .interactive_translation_pipeline import InteractiveTranslationPipeline
     from .named_entity_recognition_pipeline import NamedEntityRecognitionPipeline
     from .text_ranking_pipeline import TextRankingPipeline
     from .sentence_embedding_pipeline import SentenceEmbeddingPipeline
@@ -54,6 +55,7 @@ else:
         'feature_extraction_pipeline': ['FeatureExtractionPipeline'],
         'fill_mask_pipeline': ['FillMaskPipeline'],
         'information_extraction_pipeline': ['InformationExtractionPipeline'],
+        'interactive_translation_pipeline': ['InteractiveTranslationPipeline'],
         'named_entity_recognition_pipeline': [
             'NamedEntityRecognitionPipeline',
         ],

modelscope/pipelines/nlp/interactive_translation_pipeline.py

@@ -0,0 +1,170 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import os.path as osp
+from typing import Any, Dict
+
+import jieba
+import numpy as np
+import tensorflow as tf
+from sacremoses import MosesDetokenizer, MosesPunctNormalizer, MosesTokenizer
+from subword_nmt import apply_bpe
+
+from modelscope.metainfo import Pipelines
+from modelscope.models.base import Model
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines.base import Pipeline
+from modelscope.pipelines.builder import PIPELINES
+from modelscope.pipelines.nlp.translation_pipeline import TranslationPipeline
+from modelscope.utils.config import Config
+from modelscope.utils.constant import ModelFile, Tasks
+from modelscope.utils.logger import get_logger
+
+if tf.__version__ >= '2.0':
+    tf = tf.compat.v1
+    tf.disable_eager_execution()
+
+logger = get_logger()
+
+__all__ = ['InteractiveTranslationPipeline']
+
+
+@PIPELINES.register_module(
+    Tasks.translation, module_name=Pipelines.interactive_translation)
+class InteractiveTranslationPipeline(TranslationPipeline):
+
+    def __init__(self, model: Model, **kwargs):
+        """Build a interactive translation pipeline with a model dir or a model id in the model hub.
+
+        Args:
+            model (`str` or `Model` or module instance): A model instance or a model local dir
+                or a model id in the model hub.
+
+        Example:
+            >>> from modelscope.pipelines import pipeline
+            >>> pipeline_ins = pipeline(task=Tasks.translation,
+                model='damo/nlp_imt_translation_zh2en')
+            >>> input_sequence = 'Elon Musk, co-founder and chief executive officer of Tesla Motors.'
+            >>> input_prefix = "特斯拉汽车公司"
+            >>> print(pipeline_ins(input_sequence + "<PREFIX_SPLIT>" + input_prefix))
+        """
+        super().__init__(model=model, **kwargs)
+        model = self.model.model_dir
+        tf.reset_default_graph()
+        model_path = osp.join(
+            osp.join(model, ModelFile.TF_CHECKPOINT_FOLDER), 'ckpt-0')
+
+        self._trg_vocab = dict([
+            (w.strip(), i) for i, w in enumerate(open(self._trg_vocab_path))
+        ])
+        self._len_tgt_vocab = len(self._trg_rvocab)
+
+        self.input_wids = tf.placeholder(
+            dtype=tf.int64, shape=[None, None], name='input_wids')
+
+        self.prefix_wids = tf.placeholder(
+            dtype=tf.int64, shape=[None, None], name='prefix_wids')
+
+        self.prefix_hit = tf.placeholder(
+            dtype=tf.bool, shape=[None, None], name='prefix_hit')
+
+        self.output = {}
+
+        # preprocess
+        if self._tgt_lang == 'zh':
+            self._tgt_tok = jieba
+        else:
+            self._tgt_punct_normalizer = MosesPunctNormalizer(
+                lang=self._tgt_lang)
+            self._tgt_tok = MosesTokenizer(lang=self._tgt_lang)
+
+        # model
+        output = self.model(self.input_wids, None, self.prefix_wids,
+                            self.prefix_hit)
+        self.output.update(output)
+
+        tf_config = tf.ConfigProto(allow_soft_placement=True)
+        tf_config.gpu_options.allow_growth = True
+        self._session = tf.Session(config=tf_config)
+
+        with self._session.as_default() as sess:
+            logger.info(f'loading model from {model_path}')
+            # load model
+            model_loader = tf.train.Saver(tf.global_variables())
+            model_loader.restore(sess, model_path)
+
+    def preprocess(self, input: str) -> Dict[str, Any]:
+        input_src, prefix = input.split('<PREFIX_SPLIT>', 1)
+        if self._src_lang == 'zh':
+            input_tok = self._tok.cut(input_src)
+            input_tok = ' '.join(list(input_tok))
+        else:
+            input_src = self._punct_normalizer.normalize(input_src)
+            input_tok = self._tok.tokenize(
+                input_src, return_str=True, aggressive_dash_splits=True)
+
+        if self._tgt_lang == 'zh':
+            prefix = self._tgt_tok.lcut(prefix)
+            prefix_tok = ' '.join(list(prefix)[:-1])
+        else:
+            prefix = self._tgt_punct_normalizer.normalize(prefix)
+            prefix = self._tgt_tok.tokenize(
+                prefix, return_str=True, aggressive_dash_splits=True).split()
+            prefix_tok = ' '.join(prefix[:-1])
+
+        if len(list(prefix)) > 0:
+            subword = list(prefix)[-1]
+        else:
+            subword = ''
+
+        input_bpe = self._bpe.process_line(input_tok)
+        prefix_bpe = self._bpe.process_line(prefix_tok)
+        input_ids = np.array([[
+            self._src_vocab[w]
+            if w in self._src_vocab else self.cfg['model']['src_vocab_size']
+            for w in input_bpe.strip().split()
+        ]])
+
+        prefix_ids = np.array([[
+            self._trg_vocab[w]
+            if w in self._trg_vocab else self.cfg['model']['trg_vocab_size']
+            for w in prefix_bpe.strip().split()
+        ]])
+
+        prefix_hit = [[0] * (self._len_tgt_vocab + 1)]
+
+        if subword != '':
+            hit_state = False
+            for i, w in self._trg_rvocab.items():
+                if w.startswith(subword):
+                    prefix_hit[0][i] = 1
+                    hit_state = True
+            if hit_state is False:
+                prefix_hit = [[1] * (self._len_tgt_vocab + 1)]
+        result = {
+            'input_ids': input_ids,
+            'prefix_ids': prefix_ids,
+            'prefix_hit': np.array(prefix_hit)
+        }
+        return result
+
+    def forward(self, input: Dict[str, Any]) -> Dict[str, Any]:
+        with self._session.as_default():
+            feed_dict = {
+                self.input_wids: input['input_ids'],
+                self.prefix_wids: input['prefix_ids'],
+                self.prefix_hit: input['prefix_hit']
+            }
+            sess_outputs = self._session.run(self.output, feed_dict=feed_dict)
+            return sess_outputs
+
+    def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
+        output_seqs = inputs['output_seqs'][0]
+        wids = list(output_seqs[0]) + [0]
+        wids = wids[:wids.index(0)]
+        translation_out = ' '.join([
+            self._trg_rvocab[wid] if wid in self._trg_rvocab else '<unk>'
+            for wid in wids
+        ]).replace('@@ ', '').replace('@@', '')
+        translation_out = self._detok.detokenize(translation_out.split())
+        result = {OutputKeys.TRANSLATION: translation_out}
+        return result

tests/pipelines/test_interactive_translation_pipeline.py

@@ -0,0 +1,37 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import unittest
+
+from modelscope.pipelines import pipeline
+from modelscope.utils.constant import Tasks
+from modelscope.utils.demo_utils import DemoCompatibilityCheck
+from modelscope.utils.test_utils import test_level
+
+
+class InteractiveTranslationTest(unittest.TestCase, DemoCompatibilityCheck):
+
+    def setUp(self) -> None:
+        self.task = Tasks.translation
+        self.model_id = 'damo/nlp_imt_translation_zh2en'
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_run_with_imt_model_name_for_zh2en(self):
+        inputs = '声明补充说，沃伦的同事都深感震惊，并且希望他能够投案自首。'
+        prefix = 'The statement ad'
+        pipeline_ins = pipeline(self.task, model=self.model_id)
+        print(pipeline_ins(inputs + '<PREFIX_SPLIT>' + prefix))
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_run_with_imt_model_name_for_en2zh(self):
+        model_id = 'damo/nlp_imt_translation_en2zh'
+        inputs = 'Elon Musk, co-founder and chief executive officer of Tesla Motors.'
+        prefix = '特斯拉汽车公司'
+        pipeline_ins = pipeline(self.task, model=model_id)
+        print(pipeline_ins(inputs + '<PREFIX_SPLIT>' + prefix))
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_demo_compatibility(self):
+        self.compatibility_check()
+
+
+if __name__ == '__main__':
+    unittest.main()