From 05a407da7f8dc515d06d15c5da51f76ebaa4d527 Mon Sep 17 00:00:00 2001 From: wk258730 Date: Fri, 10 Feb 2023 08:13:18 +0000 Subject: [PATCH] add pipelines for en2zh-imt and zh2en-imt MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit *[ 交互式机器翻译-英中-通用领域-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 | 1 + modelscope/models/nlp/csanmt/translation.py | 132 +++++++++++--- modelscope/pipelines/nlp/__init__.py | 2 + .../nlp/interactive_translation_pipeline.py | 170 ++++++++++++++++++ .../test_interactive_translation_pipeline.py | 37 ++++ 5 files changed, 315 insertions(+), 27 deletions(-) create mode 100644 modelscope/pipelines/nlp/interactive_translation_pipeline.py create mode 100644 tests/pipelines/test_interactive_translation_pipeline.py diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index 11fcd4c7..9b1c1d06 100644 --- a/modelscope/metainfo.py +++ b/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' diff --git a/modelscope/models/nlp/csanmt/translation.py b/modelscope/models/nlp/csanmt/translation.py index 29d2ba84..6ece4255 100644 --- a/modelscope/models/nlp/csanmt/translation.py +++ b/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'] \ diff --git a/modelscope/pipelines/nlp/__init__.py b/modelscope/pipelines/nlp/__init__.py index 1ec4f155..0cebae80 100644 --- a/modelscope/pipelines/nlp/__init__.py +++ b/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', ], diff --git a/modelscope/pipelines/nlp/interactive_translation_pipeline.py b/modelscope/pipelines/nlp/interactive_translation_pipeline.py new file mode 100644 index 00000000..86df09bb --- /dev/null +++ b/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 + "" + 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('', 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 '' + for wid in wids + ]).replace('@@ ', '').replace('@@', '') + translation_out = self._detok.detokenize(translation_out.split()) + result = {OutputKeys.TRANSLATION: translation_out} + return result diff --git a/tests/pipelines/test_interactive_translation_pipeline.py b/tests/pipelines/test_interactive_translation_pipeline.py new file mode 100644 index 00000000..b973250a --- /dev/null +++ b/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)) + + @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)) + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_demo_compatibility(self): + self.compatibility_check() + + +if __name__ == '__main__': + unittest.main()