From 7a12181a25f6a9bebbfc32fc77379a772e44a93f Mon Sep 17 00:00:00 2001 From: "jiaqi.sjq" Date: Thu, 6 Apr 2023 19:17:07 +0800 Subject: [PATCH 1/3] [to #41669377] fix import missing Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/12240202 --- modelscope/tools/speech_tts_autolabel.py | 1 + 1 file changed, 1 insertion(+) diff --git a/modelscope/tools/speech_tts_autolabel.py b/modelscope/tools/speech_tts_autolabel.py index 0fcd41fd..a774c44f 100644 --- a/modelscope/tools/speech_tts_autolabel.py +++ b/modelscope/tools/speech_tts_autolabel.py @@ -3,6 +3,7 @@ import os import sys import zipfile +from modelscope.hub.check_model import check_local_model_is_latest from modelscope.hub.snapshot_download import snapshot_download from modelscope.utils.constant import ThirdParty from modelscope.utils.logger import get_logger From 63d5493962e4e20598d1a6973b24af64e5dafd52 Mon Sep 17 00:00:00 2001 From: "xiaoyi.zp" Date: Fri, 7 Apr 2023 14:37:18 +0800 Subject: [PATCH 2/3] add canmt translation model damo/nlp_canmt_translation_zh2en_large * add canmt translation model * add args for decode * change nlp_canmt_translation_zh2en_large damo * merge master * set default model for canmt Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/12079906 --- modelscope/metainfo.py | 11 +- modelscope/models/nlp/__init__.py | 2 + modelscope/models/nlp/canmt/__init__.py | 3 + modelscope/models/nlp/canmt/canmt_model.py | 1301 +++++++++++++++++ .../models/nlp/canmt/canmt_translation.py | 78 + .../models/nlp/canmt/sequence_generator.py | 850 +++++++++++ modelscope/pipeline_inputs.py | 2 + modelscope/pipelines/nlp/__init__.py | 2 + .../nlp/canmt_translation_pipeline.py | 91 ++ modelscope/preprocessors/__init__.py | 3 +- modelscope/preprocessors/base.py | 2 + modelscope/preprocessors/nlp/__init__.py | 4 + .../preprocessors/nlp/canmt_translation.py | 109 ++ modelscope/preprocessors/nlp/text_clean.py | 70 + modelscope/utils/constant.py | 1 + tests/pipelines/test_canmt_translation.py | 68 + 16 files changed, 2593 insertions(+), 4 deletions(-) create mode 100644 modelscope/models/nlp/canmt/__init__.py create mode 100644 modelscope/models/nlp/canmt/canmt_model.py create mode 100644 modelscope/models/nlp/canmt/canmt_translation.py create mode 100644 modelscope/models/nlp/canmt/sequence_generator.py create mode 100644 modelscope/pipelines/nlp/canmt_translation_pipeline.py create mode 100644 modelscope/preprocessors/nlp/canmt_translation.py create mode 100644 modelscope/preprocessors/nlp/text_clean.py create mode 100644 tests/pipelines/test_canmt_translation.py diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index 360f3241..9d4b111d 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -130,6 +130,7 @@ class Models(object): deberta_v2 = 'deberta_v2' veco = 'veco' translation = 'csanmt-translation' + canmt = 'canmt' space_dst = 'space-dst' space_intent = 'space-intent' space_modeling = 'space-modeling' @@ -421,6 +422,7 @@ class Pipelines(object): fill_mask = 'fill-mask' fill_mask_ponet = 'fill-mask-ponet' csanmt_translation = 'csanmt-translation' + canmt_translation = 'canmt-translation' interactive_translation = 'interactive-translation' nli = 'nli' dialog_intent_prediction = 'dialog-intent-prediction' @@ -535,6 +537,8 @@ DEFAULT_MODEL_FOR_PIPELINE = { Tasks.sentence_similarity: (Pipelines.sentence_similarity, 'damo/nlp_structbert_sentence-similarity_chinese-base'), + Tasks.competency_aware_translation: + (Pipelines.canmt_translation, 'damo/nlp_canmt_translation_zh2en_large'), Tasks.translation: (Pipelines.csanmt_translation, 'damo/nlp_csanmt_translation_zh2en'), Tasks.nli: (Pipelines.nli, 'damo/nlp_structbert_nli_chinese-base'), @@ -733,9 +737,9 @@ DEFAULT_MODEL_FOR_PIPELINE = { 'damo/nlp_structbert_faq-question-answering_chinese-base'), Tasks.crowd_counting: (Pipelines.crowd_counting, 'damo/cv_hrnet_crowd-counting_dcanet'), - Tasks.video_single_object_tracking: - (Pipelines.video_single_object_tracking, - 'damo/cv_vitb_video-single-object-tracking_ostrack'), + Tasks.video_single_object_tracking: ( + Pipelines.video_single_object_tracking, + 'damo/cv_vitb_video-single-object-tracking_ostrack'), Tasks.image_reid_person: (Pipelines.image_reid_person, 'damo/cv_passvitb_image-reid-person_market'), Tasks.text_driven_segmentation: ( @@ -993,6 +997,7 @@ class Preprocessors(object): mglm_summarization = 'mglm-summarization' sentence_piece = 'sentence-piece' translation_evaluation = 'translation-evaluation-preprocessor' + canmt_translation = 'canmt-translation' dialog_use_preprocessor = 'dialog-use-preprocessor' siamese_uie_preprocessor = 'siamese-uie-preprocessor' document_grounded_dialog_retrieval = 'document-grounded-dialog-retrieval' diff --git a/modelscope/models/nlp/__init__.py b/modelscope/models/nlp/__init__.py index 66a53a00..831c0f9c 100644 --- a/modelscope/models/nlp/__init__.py +++ b/modelscope/models/nlp/__init__.py @@ -19,6 +19,7 @@ if TYPE_CHECKING: from .bloom import BloomModel from .codegeex import CodeGeeXForCodeTranslation, CodeGeeXForCodeGeneration from .csanmt import CsanmtForTranslation + from .canmt import CanmtForTranslation from .deberta_v2 import DebertaV2ForMaskedLM, DebertaV2Model from .gpt_neo import GPTNeoModel from .gpt2 import GPT2Model @@ -87,6 +88,7 @@ else: ], 'bloom': ['BloomModel'], 'csanmt': ['CsanmtForTranslation'], + 'canmt': ['CanmtForTranslation'], 'codegeex': ['CodeGeeXForCodeTranslation', 'CodeGeeXForCodeGeneration'], 'deberta_v2': ['DebertaV2ForMaskedLM', 'DebertaV2Model'], diff --git a/modelscope/models/nlp/canmt/__init__.py b/modelscope/models/nlp/canmt/__init__.py new file mode 100644 index 00000000..affe5b1f --- /dev/null +++ b/modelscope/models/nlp/canmt/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +from .canmt_translation import CanmtForTranslation diff --git a/modelscope/models/nlp/canmt/canmt_model.py b/modelscope/models/nlp/canmt/canmt_model.py new file mode 100644 index 00000000..e55c6f89 --- /dev/null +++ b/modelscope/models/nlp/canmt/canmt_model.py @@ -0,0 +1,1301 @@ +# Part of the implementation is borrowed and modified from FAIRSEQ, +# publicly available at https://github.com/facebookresearch/fairseq +# Copyright 2022-2023 The Alibaba MT Team Authors. All rights reserved. +import math +from typing import Any, Dict, List, Optional, Tuple + +import numpy +import torch +import torch.nn as nn +from fairseq import utils +from fairseq.distributed import fsdp_wrap +from fairseq.models import (FairseqEncoder, FairseqEncoderDecoderModel, + FairseqIncrementalDecoder, register_model, + register_model_architecture) +from fairseq.modules import (AdaptiveSoftmax, BaseLayer, FairseqDropout, + LayerDropModuleList, LayerNorm, + PositionalEmbedding, + SinusoidalPositionalEmbedding, + TransformerDecoderLayer, TransformerEncoderLayer) +from fairseq.modules.checkpoint_activations import checkpoint_wrapper +from fairseq.modules.quant_noise import quant_noise as apply_quant_noise_ +from torch import Tensor + +DEFAULT_MAX_SOURCE_POSITIONS = 1024 +DEFAULT_MAX_TARGET_POSITIONS = 1024 + +DEFAULT_MIN_PARAMS_TO_WRAP = int(1e8) + + +class CanmtModel(FairseqEncoderDecoderModel): + """ + + Args: + encoder (TransformerEncoder): the encoder + decoder (TransformerDecoder): the decoder + + The CanmtModel provides the following named architectures and + command-line arguments: + + .. argparse:: + :ref: fairseq.models.transformer_parser + :prog: + """ + + def __init__(self, args, encoder, decoder, second_decoder): + super().__init__(encoder, decoder) + self.args = args + self.supports_align_args = True + self.encoder = encoder + self.decoder = decoder + self.second_decoder = second_decoder + + @staticmethod + def add_args(parser): + """Add model-specific arguments to the parser.""" + # fmt: off + parser.add_argument( + '--activation-fn', + choices=utils.get_available_activation_fns(), + help='activation function to use') + parser.add_argument( + '--dropout', type=float, metavar='D', help='dropout probability') + parser.add_argument( + '--attention-dropout', + type=float, + metavar='D', + help='dropout probability for attention weights') + parser.add_argument( + '--activation-dropout', + '--relu-dropout', + type=float, + metavar='D', + help='dropout probability after activation in FFN.') + parser.add_argument( + '--encoder-embed-path', + type=str, + metavar='STR', + help='path to pre-trained encoder embedding') + parser.add_argument( + '--encoder-embed-dim', + type=int, + metavar='N', + help='encoder embedding dimension') + parser.add_argument( + '--encoder-ffn-embed-dim', + type=int, + metavar='N', + help='encoder embedding dimension for FFN') + parser.add_argument( + '--encoder-layers', + type=int, + metavar='N', + help='num encoder layers') + parser.add_argument( + '--encoder-attention-heads', + type=int, + metavar='N', + help='num encoder attention heads') + parser.add_argument( + '--encoder-normalize-before', + action='store_true', + help='apply layernorm before each encoder block') + parser.add_argument( + '--encoder-learned-pos', + action='store_true', + help='use learned positional embeddings in the encoder') + parser.add_argument( + '--decoder-embed-path', + type=str, + metavar='STR', + help='path to pre-trained decoder embedding') + parser.add_argument( + '--decoder-embed-dim', + type=int, + metavar='N', + help='decoder embedding dimension') + parser.add_argument( + '--decoder-ffn-embed-dim', + type=int, + metavar='N', + help='decoder embedding dimension for FFN') + parser.add_argument( + '--decoder-layers', + type=int, + metavar='N', + help='num decoder layers') + parser.add_argument( + '--decoder-attention-heads', + type=int, + metavar='N', + help='num decoder attention heads') + parser.add_argument( + '--decoder-learned-pos', + action='store_true', + help='use learned positional embeddings in the decoder') + parser.add_argument( + '--decoder-normalize-before', + action='store_true', + help='apply layernorm before each decoder block') + parser.add_argument( + '--decoder-output-dim', + type=int, + metavar='N', + help='decoder output dimension (extra linear layer ' + 'if different from decoder embed dim') + parser.add_argument( + '--share-decoder-input-output-embed', + action='store_true', + help='share decoder input and output embeddings') + parser.add_argument( + '--share-all-embeddings', + action='store_true', + help='share encoder, decoder and output embeddings' + ' (requires shared dictionary and embed dim)') + parser.add_argument( + '--no-token-positional-embeddings', + default=False, + action='store_true', + help= + 'if set, disables positional embeddings (outside self attention)') + parser.add_argument( + '--adaptive-softmax-cutoff', + metavar='EXPR', + help='comma separated list of adaptive softmax cutoff points. ' + 'Must be used with adaptive_loss criterion'), + parser.add_argument( + '--adaptive-softmax-dropout', + type=float, + metavar='D', + help='sets adaptive softmax dropout for the tail projections') + parser.add_argument( + '--layernorm-embedding', + action='store_true', + help='add layernorm to embedding') + parser.add_argument( + '--no-scale-embedding', + action='store_true', + help='if True, dont scale embeddings') + parser.add_argument( + '--checkpoint-activations', + action='store_true', + help='checkpoint activations at each layer, which saves GPU ' + 'memory usage at the cost of some additional compute') + parser.add_argument( + '--offload-activations', + action='store_true', + help='checkpoint activations at each layer, then save to gpu.' + 'Sets --checkpoint-activations.') + parser.add_argument( + '--no-cross-attention', + default=False, + action='store_true', + help='do not perform cross-attention') + parser.add_argument( + '--cross-self-attention', + default=False, + action='store_true', + help='perform cross+self-attention') + parser.add_argument( + '--encoder-layerdrop', + type=float, + metavar='D', + default=0, + help='LayerDrop probability for encoder') + parser.add_argument( + '--decoder-layerdrop', + type=float, + metavar='D', + default=0, + help='LayerDrop probability for decoder') + parser.add_argument( + '--encoder-layers-to-keep', + default=None, + help='which layers to *keep* when pruning as a comma-separated list' + ) + parser.add_argument( + '--decoder-layers-to-keep', + default=None, + help='which layers to *keep* when pruning as a comma-separated list' + ) + parser.add_argument( + '--quant-noise-pq', + type=float, + metavar='D', + default=0, + help='iterative PQ quantization noise at training time') + parser.add_argument( + '--quant-noise-pq-block-size', + type=int, + metavar='D', + default=8, + help='block size of quantization noise at training time') + parser.add_argument( + '--quant-noise-scalar', + type=float, + metavar='D', + default=0, + help= + 'scalar quantization noise and scalar quantization at training time' + ) + parser.add_argument( + '--min-params-to-wrap', + type=int, + metavar='D', + default=DEFAULT_MIN_PARAMS_TO_WRAP, + help= + ('minimum number of params for a layer to be wrapped with FSDP() when ' + 'training with --ddp-backend=fully_sharded. Smaller values will ' + 'improve memory efficiency, but may make torch.distributed ' + 'communication less efficient due to smaller input sizes. This option ' + 'is set to 0 (i.e., always wrap) when --checkpoint-activations or ' + '--offload-activations are passed.')) + # fmt: on + + @classmethod + def build_model(cls, args, task): + """Build a new model instance.""" + + # make sure all arguments are present in older models + base_architecture(args) + + if args.encoder_layers_to_keep: + args.encoder_layers = len(args.encoder_layers_to_keep.split(',')) + if args.decoder_layers_to_keep: + args.decoder_layers = len(args.decoder_layers_to_keep.split(',')) + + if getattr(args, 'max_source_positions', None) is None: + args.max_source_positions = DEFAULT_MAX_SOURCE_POSITIONS + if getattr(args, 'max_target_positions', None) is None: + args.max_target_positions = DEFAULT_MAX_TARGET_POSITIONS + + src_dict, tgt_dict = task.vocab_src, task.vocab_tgt + + if args.share_all_embeddings: + if src_dict != tgt_dict: + raise ValueError( + '--share-all-embeddings requires a joined dictionary') + if args.encoder_embed_dim != args.decoder_embed_dim: + raise ValueError( + '--share-all-embeddings requires --encoder-embed-dim to match --decoder-embed-dim' + ) + if args.decoder_embed_path and \ + (args.decoder_embed_path != args.encoder_embed_path): + raise ValueError( + '--share-all-embeddings not compatible with --decoder-embed-path' + ) + encoder_embed_tokens = cls.build_embedding(args, src_dict, + args.encoder_embed_dim, + args.encoder_embed_path) + decoder_embed_tokens = encoder_embed_tokens + args.share_decoder_input_output_embed = True + else: + encoder_embed_tokens = cls.build_embedding(args, src_dict, + args.encoder_embed_dim, + args.encoder_embed_path) + decoder_embed_tokens = cls.build_embedding(args, tgt_dict, + args.decoder_embed_dim, + args.decoder_embed_path) + if getattr(args, 'offload_activations', False): + args.checkpoint_activations = True # offloading implies checkpointing + encoder = cls.build_encoder(args, src_dict, encoder_embed_tokens) + decoder = cls.build_decoder(args, tgt_dict, decoder_embed_tokens) + second_decoder = cls.build_decoder(args, src_dict, + encoder_embed_tokens) + if not args.share_all_embeddings: + min_params_to_wrap = getattr(args, 'min_params_to_wrap', + DEFAULT_MIN_PARAMS_TO_WRAP) + # fsdp_wrap is a no-op when --ddp-backend != fully_sharded + encoder = fsdp_wrap(encoder, min_num_params=min_params_to_wrap) + decoder = fsdp_wrap(decoder, min_num_params=min_params_to_wrap) + return cls(args, encoder, decoder, second_decoder) + + @classmethod + def build_embedding(cls, args, dictionary, embed_dim, path=None): + num_embeddings = len(dictionary) + padding_idx = dictionary.pad() + + emb = Embedding(num_embeddings, embed_dim, padding_idx) + # if provided, load from preloaded dictionaries + if path: + embed_dict = utils.parse_embedding(path) + utils.load_embedding(embed_dict, dictionary, emb) + return emb + + @classmethod + def build_encoder(cls, args, src_dict, embed_tokens): + return TransformerEncoder(args, src_dict, embed_tokens) + + @classmethod + def build_decoder(cls, args, tgt_dict, embed_tokens): + return TransformerDecoder( + args, + tgt_dict, + embed_tokens, + no_encoder_attn=getattr(args, 'no_cross_attention', False), + ) + + def forward( + self, + src_tokens, + src_lengths, + prev_output_tokens, + prev_src_tokens, + return_all_hiddens: bool = True, + features_only: bool = False, + alignment_layer: Optional[int] = None, + alignment_heads: Optional[int] = None, + ): + """ + Run the forward pass for an encoder-decoder model. + + Copied from the base class, but without ``**kwargs``, + which are not supported by TorchScript. + """ + + encoder_out = self.encoder( + src_tokens, + src_lengths=src_lengths, + return_all_hiddens=return_all_hiddens) + decoder_out = self.decoder( + prev_output_tokens, + encoder_out=encoder_out, + features_only=features_only, + alignment_layer=alignment_layer, + alignment_heads=alignment_heads, + src_lengths=src_lengths, + return_all_hiddens=return_all_hiddens, + ) + + decoder_out_re = self.decoder( + prev_output_tokens, + encoder_out=None, + features_only=features_only, + full_context_alignment=True, + alignment_layer=alignment_layer, + alignment_heads=alignment_heads, + src_lengths=src_lengths, + return_all_hiddens=return_all_hiddens, + ) + decoder_out_tensor = decoder_out_re[1]['last_layer'] + decoder_padding = decoder_out_re[1]['self_attn_padding_mask'] + + decoder_kvs = { + 'encoder_out': [decoder_out_tensor], + 'encoder_padding_mask': [decoder_padding] + } + src_out = self.second_decoder( + prev_src_tokens, + encoder_out=decoder_kvs, + features_only=features_only, + alignment_layer=alignment_layer, + alignment_heads=alignment_heads, + src_lengths=None, + return_all_hiddens=return_all_hiddens, + ) + return decoder_out, src_out, decoder_kvs + + @torch.jit.export + def get_normalized_probs( + self, + net_output: Tuple[Tensor, Optional[Dict[str, List[Optional[Tensor]]]]], + log_probs: bool, + sample: Optional[Dict[str, Tensor]] = None, + ): + """Get normalized probabilities (or log probs) from a net's output.""" + return self.get_normalized_probs_scriptable(net_output, log_probs, + sample) + + def forward_decoder( + self, + tokens, + encoder_outs: Dict[str, List[Tensor]], + incremental_states: Dict[str, Dict[str, Optional[Tensor]]], + temperature: float = 1.0, + ): + encoder_out: Optional[Dict[str, List[Tensor]]] = None + encoder_out = encoder_outs + # decode + decoder_out = self.decoder.forward( + tokens, + encoder_out=encoder_out, + incremental_state=incremental_states, + ) + + attn: Optional[Tensor] = None + decoder_len = len(decoder_out) + if decoder_len > 1 and decoder_out[1] is not None: + if isinstance(decoder_out[1], Tensor): + attn = decoder_out[1] + else: + attn_holder = decoder_out[1]['attn'] + if isinstance(attn_holder, Tensor): + attn = attn_holder + elif attn_holder is not None: + attn = attn_holder[0] + if attn is not None: + attn = attn[:, -1, :] + + decoder_out_tuple = ( + decoder_out[0][:, -1:, :].div_(temperature), + None if decoder_len <= 1 else decoder_out[1], + ) + probs = self.get_normalized_probs( + decoder_out_tuple, log_probs=True, sample=None) + probs = probs[:, -1, :] + decoder_out_tensor = decoder_out[1]['last_layer'] + return probs, attn, decoder_out_tensor + + def forward_decoder_src( + self, + tokens, + encoder_outs: Dict[str, List[Tensor]], + incremental_states: Dict[str, Dict[str, Optional[Tensor]]], + temperature: float = 1.0, + ): + encoder_out: Optional[Dict[str, List[Tensor]]] = None + encoder_out = encoder_outs + # decode each model + decoder_out = self.second_decoder.forward( + tokens, encoder_out=encoder_out) + + attn: Optional[Tensor] = None + decoder_len = len(decoder_out) + if decoder_len > 1 and decoder_out[1] is not None: + if isinstance(decoder_out[1], Tensor): + attn = decoder_out[1] + else: + attn_holder = decoder_out[1]['attn'] + if isinstance(attn_holder, Tensor): + attn = attn_holder + elif attn_holder is not None: + attn = attn_holder[0] + if attn is not None: + attn = attn[:, -1, :] + + decoder_out_tuple = ( + decoder_out[0][:, -1:, :].div_(temperature), + None if decoder_len <= 1 else decoder_out[1], + ) + probs = self.get_normalized_probs( + decoder_out_tuple, log_probs=True, sample=None) + probs = probs[:, -1, :] + decoder_out_tensor = decoder_out[1]['last_layer'] + return probs, attn, decoder_out_tensor, decoder_out + + def forward_encoder(self, net_input: Dict[str, Tensor]): + encoder_input = { + k: v + for k, v in net_input.items() if k != 'prev_output_tokens' + and k != 'prev_src_tokens' and k != 'sources' + } + return self.encoder.forward_torchscript(encoder_input) + + def reorder_encoder_out(self, encoder_outs: Optional[Dict[str, + List[Tensor]]], + new_order): + """ + Reorder encoder output according to *new_order*. + + Args: + encoder_out: output from the ``forward()`` method + new_order (LongTensor): desired order + + Returns: + *encoder_out* rearranged according to *new_order* + """ + assert encoder_outs is not None + return self.encoder.reorder_encoder_out(encoder_outs, new_order) + + def reorder_incremental_state( + self, + incremental_states: Dict[str, Dict[str, Optional[Tensor]]], + new_order, + ): + self.decoder.reorder_incremental_state_scripting( + incremental_states, new_order) + + +class TransformerEncoder(FairseqEncoder): + """ + Transformer encoder consisting of *args.encoder_layers* layers. Each layer + is a :class:`TransformerEncoderLayer`. + + Args: + args (argparse.Namespace): parsed command-line arguments + dictionary (~fairseq.data.Dictionary): encoding dictionary + embed_tokens (torch.nn.Embedding): input embedding + """ + + def __init__(self, args, dictionary, embed_tokens): + self.args = args + super().__init__(dictionary) + self.register_buffer('version', torch.Tensor([3])) + + self.dropout_module = FairseqDropout( + args.dropout, module_name=self.__class__.__name__) + self.encoder_layerdrop = args.encoder_layerdrop + + embed_dim = embed_tokens.embedding_dim + self.padding_idx = embed_tokens.padding_idx + self.max_source_positions = args.max_source_positions + self.embed_tokens = embed_tokens + + self.embed_scale = 1.0 if args.no_scale_embedding else math.sqrt( + embed_dim) + + self.embed_positions = ( + PositionalEmbedding( + args.max_source_positions, + embed_dim, + self.padding_idx, + learned=args.encoder_learned_pos, + ) if not args.no_token_positional_embeddings else None) + export = getattr(args, 'export', False) + if getattr(args, 'layernorm_embedding', False): + self.layernorm_embedding = LayerNorm(embed_dim, export=export) + else: + self.layernorm_embedding = None + + if not args.adaptive_input and args.quant_noise_pq > 0: + self.quant_noise = apply_quant_noise_( + nn.Linear(embed_dim, embed_dim, bias=False), + args.quant_noise_pq, + args.quant_noise_pq_block_size, + ) + else: + self.quant_noise = None + + if self.encoder_layerdrop > 0.0: + self.layers = LayerDropModuleList(p=self.encoder_layerdrop) + else: + self.layers = nn.ModuleList([]) + self.layers.extend([ + self.build_encoder_layer(args) for i in range(args.encoder_layers) + ]) + self.num_layers = len(self.layers) + + if args.encoder_normalize_before: + self.layer_norm = LayerNorm(embed_dim, export=export) + else: + self.layer_norm = None + + def build_encoder_layer(self, args): + layer = TransformerEncoderLayer(args) + checkpoint = getattr(args, 'checkpoint_activations', False) + if checkpoint: + offload_to_cpu = getattr(args, 'offload_activations', False) + layer = checkpoint_wrapper(layer, offload_to_cpu=offload_to_cpu) + min_params_to_wrap = ( + getattr(args, 'min_params_to_wrap', DEFAULT_MIN_PARAMS_TO_WRAP) + if not checkpoint else 0) + layer = fsdp_wrap(layer, min_num_params=min_params_to_wrap) + return layer + + def forward_embedding(self, + src_tokens, + token_embedding: Optional[torch.Tensor] = None): + # embed tokens and positions + if token_embedding is None: + token_embedding = self.embed_tokens(src_tokens) + x = embed = self.embed_scale * token_embedding + if self.embed_positions is not None: + x = embed + self.embed_positions(src_tokens) + if self.layernorm_embedding is not None: + x = self.layernorm_embedding(x) + x = self.dropout_module(x) + if self.quant_noise is not None: + x = self.quant_noise(x) + return x, embed + + def forward( + self, + src_tokens, + src_lengths: Optional[torch.Tensor] = None, + return_all_hiddens: bool = False, + token_embeddings: Optional[torch.Tensor] = None, + ): + """ + Args: + src_tokens (LongTensor): tokens in the source language of shape + `(batch, src_len)` + src_lengths (torch.LongTensor): lengths of each source sentence of + shape `(batch)` + return_all_hiddens (bool, optional): also return all of the + intermediate hidden states (default: False). + token_embeddings (torch.Tensor, optional): precomputed embeddings + default `None` will recompute embeddings + + Returns: + dict: + - **encoder_out** (Tensor): the last encoder layer's output of + shape `(src_len, batch, embed_dim)` + - **encoder_padding_mask** (ByteTensor): the positions of + padding elements of shape `(batch, src_len)` + - **encoder_embedding** (Tensor): the (scaled) embedding lookup + of shape `(batch, src_len, embed_dim)` + - **encoder_states** (List[Tensor]): all intermediate + hidden states of shape `(src_len, batch, embed_dim)`. + Only populated if *return_all_hiddens* is True. + """ + return self.forward_scriptable(src_tokens, src_lengths, + return_all_hiddens, token_embeddings) + + def forward_scriptable( + self, + src_tokens, + src_lengths: Optional[torch.Tensor] = None, + return_all_hiddens: bool = False, + token_embeddings: Optional[torch.Tensor] = None, + ): + """ + Args: + src_tokens (LongTensor): tokens in the source language of shape + `(batch, src_len)` + src_lengths (torch.LongTensor): lengths of each source sentence of + shape `(batch)` + return_all_hiddens (bool, optional): also return all of the + intermediate hidden states (default: False). + token_embeddings (torch.Tensor, optional): precomputed embeddings + default `None` will recompute embeddings + + Returns: + dict: + - **encoder_out** (Tensor): the last encoder layer's output of + shape `(src_len, batch, embed_dim)` + - **encoder_padding_mask** (ByteTensor): the positions of + padding elements of shape `(batch, src_len)` + - **encoder_embedding** (Tensor): the (scaled) embedding lookup + of shape `(batch, src_len, embed_dim)` + - **encoder_states** (List[Tensor]): all intermediate + hidden states of shape `(src_len, batch, embed_dim)`. + Only populated if *return_all_hiddens* is True. + """ + # compute padding mask + encoder_padding_mask = src_tokens.eq(self.padding_idx) + has_pads = src_tokens.device.type == 'xla' or encoder_padding_mask.any( + ) + x, encoder_embedding = self.forward_embedding(src_tokens, + token_embeddings) + + # account for padding while computing the representation + if has_pads: + x = x * (1 - encoder_padding_mask.unsqueeze(-1).type_as(x)) + + # B x T x C -> T x B x C + x = x.transpose(0, 1) + + encoder_states = [] + + if return_all_hiddens: + encoder_states.append(x) + + # encoder layers + for layer in self.layers: + x = layer( + x, + encoder_padding_mask=encoder_padding_mask + if has_pads else None) + if return_all_hiddens: + assert encoder_states is not None + encoder_states.append(x) + + if self.layer_norm is not None: + x = self.layer_norm(x) + + return { + 'encoder_out': [x], # T x B x C + 'encoder_padding_mask': [encoder_padding_mask], # B x T + 'encoder_embedding': [encoder_embedding], # B x T x C + 'encoder_states': encoder_states, # List[T x B x C] + 'src_tokens': [], + 'src_lengths': [], + } + + @torch.jit.export + def reorder_encoder_out(self, encoder_out: Dict[str, List[Tensor]], + new_order): + """ + Reorder encoder output according to *new_order*. + + Args: + encoder_out: output from the ``forward()`` method + new_order (LongTensor): desired order + + Returns: + *encoder_out* rearranged according to *new_order* + """ + if len(encoder_out['encoder_out']) == 0: + new_encoder_out = [] + else: + new_encoder_out = [ + encoder_out['encoder_out'][0].index_select(1, new_order) + ] + if len(encoder_out['encoder_padding_mask']) == 0: + new_encoder_padding_mask = [] + else: + new_encoder_padding_mask = [ + encoder_out['encoder_padding_mask'][0].index_select( + 0, new_order) + ] + if len(encoder_out['encoder_embedding']) == 0: + new_encoder_embedding = [] + else: + new_encoder_embedding = [ + encoder_out['encoder_embedding'][0].index_select(0, new_order) + ] + + if len(encoder_out['src_tokens']) == 0: + src_tokens = [] + else: + src_tokens = [ + (encoder_out['src_tokens'][0]).index_select(0, new_order) + ] + + if len(encoder_out['src_lengths']) == 0: + src_lengths = [] + else: + src_lengths = [ + (encoder_out['src_lengths'][0]).index_select(0, new_order) + ] + + encoder_states = encoder_out['encoder_states'] + if len(encoder_states) > 0: + for idx, state in enumerate(encoder_states): + encoder_states[idx] = state.index_select(1, new_order) + + return { + 'encoder_out': new_encoder_out, # T x B x C + 'encoder_padding_mask': new_encoder_padding_mask, # B x T + 'encoder_embedding': new_encoder_embedding, # B x T x C + 'encoder_states': encoder_states, # List[T x B x C] + 'src_tokens': src_tokens, # B x T + 'src_lengths': src_lengths, # B x 1 + } + + def max_positions(self): + """Maximum input length supported by the encoder.""" + if self.embed_positions is None: + return self.max_source_positions + return min(self.max_source_positions, + self.embed_positions.max_positions) + + def upgrade_state_dict_named(self, state_dict, name): + """Upgrade a (possibly old) state dict for new versions of fairseq.""" + if isinstance(self.embed_positions, SinusoidalPositionalEmbedding): + weights_key = '{}.embed_positions.weights'.format(name) + if weights_key in state_dict: + print('deleting {0}'.format(weights_key)) + del state_dict[weights_key] + state_dict['{}.embed_positions._float_tensor'.format( + name)] = torch.FloatTensor(1) + for i in range(self.num_layers): + # update layer norms + self.layers[i].upgrade_state_dict_named( + state_dict, '{}.layers.{}'.format(name, i)) + + version_key = '{}.version'.format(name) + if utils.item(state_dict.get(version_key, torch.Tensor([1]))[0]) < 2: + # earlier checkpoints did not normalize after the stack of layers + self.layer_norm = None + self.normalize = False + state_dict[version_key] = torch.Tensor([1]) + return state_dict + + +class TransformerDecoder(FairseqIncrementalDecoder): + """ + Transformer decoder consisting of *args.decoder_layers* layers. Each layer + is a :class:`TransformerDecoderLayer`. + + Args: + args (argparse.Namespace): parsed command-line arguments + dictionary (~fairseq.data.Dictionary): decoding dictionary + embed_tokens (torch.nn.Embedding): output embedding + no_encoder_attn (bool, optional): whether to attend to encoder outputs + (default: False). + """ + + def __init__( + self, + args, + dictionary, + embed_tokens, + no_encoder_attn=False, + output_projection=None, + ): + self.args = args + super().__init__(dictionary) + self.register_buffer('version', torch.Tensor([3])) + self._future_mask = torch.empty(0) + + self.dropout_module = FairseqDropout( + args.dropout, module_name=self.__class__.__name__) + self.decoder_layerdrop = args.decoder_layerdrop + self.share_input_output_embed = args.share_decoder_input_output_embed + + input_embed_dim = embed_tokens.embedding_dim + embed_dim = args.decoder_embed_dim + self.embed_dim = embed_dim + self.output_embed_dim = args.decoder_output_dim + + self.padding_idx = embed_tokens.padding_idx + self.max_target_positions = args.max_target_positions + + self.embed_tokens = embed_tokens + + self.embed_scale = 1.0 if args.no_scale_embedding else math.sqrt( + embed_dim) + + if not args.adaptive_input and args.quant_noise_pq > 0: + self.quant_noise = apply_quant_noise_( + nn.Linear(embed_dim, embed_dim, bias=False), + args.quant_noise_pq, + args.quant_noise_pq_block_size, + ) + else: + self.quant_noise = None + + self.project_in_dim = ( + Linear(input_embed_dim, embed_dim, bias=False) + if embed_dim != input_embed_dim else None) + self.embed_positions = ( + PositionalEmbedding( + self.max_target_positions, + embed_dim, + self.padding_idx, + learned=args.decoder_learned_pos, + ) if not args.no_token_positional_embeddings else None) + export = getattr(args, 'export', False) + if getattr(args, 'layernorm_embedding', False): + self.layernorm_embedding = LayerNorm(embed_dim, export=export) + else: + self.layernorm_embedding = None + + self.cross_self_attention = getattr(args, 'cross_self_attention', + False) + + if self.decoder_layerdrop > 0.0: + self.layers = LayerDropModuleList(p=self.decoder_layerdrop) + else: + self.layers = nn.ModuleList([]) + self.layers.extend([ + self.build_decoder_layer(args, no_encoder_attn) + for _ in range(args.decoder_layers) + ]) + self.num_layers = len(self.layers) + + if args.decoder_normalize_before and not getattr( + args, 'no_decoder_final_norm', False): + self.layer_norm = LayerNorm(embed_dim, export=export) + else: + self.layer_norm = None + + self.project_out_dim = ( + Linear(embed_dim, self.output_embed_dim, bias=False) + if embed_dim != self.output_embed_dim + and not args.tie_adaptive_weights else None) + + self.adaptive_softmax = None + self.output_projection = output_projection + if self.output_projection is None: + self.build_output_projection(args, dictionary, embed_tokens) + + def build_output_projection(self, args, dictionary, embed_tokens): + if args.adaptive_softmax_cutoff is not None: + self.adaptive_softmax = AdaptiveSoftmax( + len(dictionary), + self.output_embed_dim, + utils.eval_str_list(args.adaptive_softmax_cutoff, type=int), + dropout=args.adaptive_softmax_dropout, + adaptive_inputs=embed_tokens + if args.tie_adaptive_weights else None, + factor=args.adaptive_softmax_factor, + tie_proj=args.tie_adaptive_proj, + ) + elif self.share_input_output_embed: + self.output_projection = nn.Linear( + self.embed_tokens.weight.shape[1], + self.embed_tokens.weight.shape[0], + bias=False, + ) + self.output_projection.weight = self.embed_tokens.weight + else: + self.output_projection = nn.Linear( + self.output_embed_dim, len(dictionary), bias=False) + nn.init.normal_( + self.output_projection.weight, + mean=0, + std=self.output_embed_dim**-0.5) + num_base_layers = getattr(args, 'base_layers', 0) + for i in range(num_base_layers): + self.layers.insert( + ((i + 1) * args.decoder_layers) // (num_base_layers + 1), + BaseLayer(args), + ) + + def build_decoder_layer(self, args, no_encoder_attn=False): + layer = TransformerDecoderLayer(args, no_encoder_attn) + checkpoint = getattr(args, 'checkpoint_activations', False) + if checkpoint: + offload_to_cpu = getattr(args, 'offload_activations', False) + layer = checkpoint_wrapper(layer, offload_to_cpu=offload_to_cpu) + min_params_to_wrap = ( + getattr(args, 'min_params_to_wrap', DEFAULT_MIN_PARAMS_TO_WRAP) + if not checkpoint else 0) + layer = fsdp_wrap(layer, min_num_params=min_params_to_wrap) + return layer + + def forward( + self, + prev_output_tokens, + encoder_out: Optional[Dict[str, List[Tensor]]] = None, + incremental_state: Optional[Dict[str, Dict[str, + Optional[Tensor]]]] = None, + features_only: bool = False, + full_context_alignment: bool = False, + alignment_layer: Optional[int] = None, + alignment_heads: Optional[int] = None, + src_lengths: Optional[Any] = None, + return_all_hiddens: bool = False, + ): + """ + Args: + prev_output_tokens (LongTensor): previous decoder outputs of shape + `(batch, tgt_len)`, for teacher forcing + encoder_out (optional): output from the encoder, used for + encoder-side attention, should be of size T x B x C + incremental_state (dict): dictionary used for storing state during + :ref:`Incremental decoding` + features_only (bool, optional): only return features without + applying output layer (default: False). + full_context_alignment (bool, optional): don't apply + auto-regressive mask to self-attention (default: False). + + Returns: + tuple: + - the decoder's output of shape `(batch, tgt_len, vocab)` + - a dictionary with any model-specific outputs + """ + + x, extra = self.extract_features( + prev_output_tokens, + encoder_out=encoder_out, + incremental_state=incremental_state, + full_context_alignment=full_context_alignment, + alignment_layer=alignment_layer, + alignment_heads=alignment_heads, + ) + + if not features_only: + x = self.output_layer(x) + return x, extra + + def extract_features( + self, + prev_output_tokens, + encoder_out: Optional[Dict[str, List[Tensor]]], + incremental_state: Optional[Dict[str, Dict[str, + Optional[Tensor]]]] = None, + full_context_alignment: bool = False, + alignment_layer: Optional[int] = None, + alignment_heads: Optional[int] = None, + ): + return self.extract_features_scriptable( + prev_output_tokens, + encoder_out, + incremental_state, + full_context_alignment, + alignment_layer, + alignment_heads, + ) + + """ + A scriptable subclass of this class has an extract_features method and calls + super().extract_features, but super() is not supported in torchscript. A copy of + this function is made to be used in the subclass instead. + """ + + def extract_features_scriptable( + self, + prev_output_tokens, + encoder_out: Optional[Dict[str, List[Tensor]]], + incremental_state: Optional[Dict[str, Dict[str, + Optional[Tensor]]]] = None, + full_context_alignment: bool = False, + alignment_layer: Optional[int] = None, + alignment_heads: Optional[int] = None, + ): + """ + Similar to *forward* but only return features. + + Includes several features from "Jointly Learning to Align and + Translate with Transformer Models" (Garg et al., EMNLP 2019). + + Args: + full_context_alignment (bool, optional): don't apply + auto-regressive mask to self-attention (default: False). + alignment_layer (int, optional): return mean alignment over + heads at this layer (default: last layer). + alignment_heads (int, optional): only average alignment over + this many heads (default: all heads). + + Returns: + tuple: + - the decoder's features of shape `(batch, tgt_len, embed_dim)` + - a dictionary with any model-specific outputs + """ + bs, slen = prev_output_tokens.size() + if alignment_layer is None: + alignment_layer = self.num_layers - 1 + + enc: Optional[Tensor] = None + padding_mask: Optional[Tensor] = None + if encoder_out is not None and len(encoder_out['encoder_out']) > 0: + enc = encoder_out['encoder_out'][0] + assert (enc.size()[1] == bs + ), f'Expected enc.shape == (t, {bs}, c) got {enc.shape}' + if encoder_out is not None and len( + encoder_out['encoder_padding_mask']) > 0: + padding_mask = encoder_out['encoder_padding_mask'][0] + + # embed positions + positions = None + if self.embed_positions is not None: + positions = self.embed_positions( + prev_output_tokens, incremental_state=incremental_state) + if incremental_state is not None: + prev_output_tokens = prev_output_tokens[:, -1:] + if positions is not None: + positions = positions[:, -1:] + + # embed tokens and positions + x = self.embed_scale * self.embed_tokens(prev_output_tokens) + + if self.quant_noise is not None: + x = self.quant_noise(x) + + if self.project_in_dim is not None: + x = self.project_in_dim(x) + + if positions is not None: + x += positions + + if self.layernorm_embedding is not None: + x = self.layernorm_embedding(x) + + x = self.dropout_module(x) + + # B x T x C -> T x B x C + x = x.transpose(0, 1) + + self_attn_padding_mask: Optional[Tensor] = None + if self.cross_self_attention or prev_output_tokens.eq( + self.padding_idx).any(): + self_attn_padding_mask = prev_output_tokens.eq(self.padding_idx) + + # decoder layers + attn: Optional[Tensor] = None + inner_states: List[Optional[Tensor]] = [x] + for idx, layer in enumerate(self.layers): + if incremental_state is None and not full_context_alignment: + self_attn_mask = self.buffered_future_mask(x) + else: + self_attn_mask = None + + x, layer_attn, self_attn_hidden = layer( + x, + enc, + padding_mask, + incremental_state, + self_attn_mask=self_attn_mask, + self_attn_padding_mask=self_attn_padding_mask, + need_attn=bool((idx == alignment_layer)), + need_head_weights=bool((idx == alignment_layer)), + ) + inner_states.append(x) + if layer_attn is not None and idx == alignment_layer: + attn = layer_attn.float().to(x) + + if attn is not None: + if alignment_heads is not None: + attn = attn[:alignment_heads] + + attn = attn.mean(dim=0) + + if self.layer_norm is not None: + x = self.layer_norm(x) + last_layer = x + # T x B x C -> B x T x C + x = x.transpose(0, 1) + if self.project_out_dim is not None: + x = self.project_out_dim(x) + return x, { + 'attn': [attn], + 'inner_states': inner_states, + 'last_layer': last_layer, + 'self_attn_padding_mask': self_attn_padding_mask + } + + def output_layer(self, features): + """Project features to the vocabulary size.""" + if self.adaptive_softmax is None: + # project back to size of vocabulary + return self.output_projection(features) + else: + return features + + def max_positions(self): + """Maximum output length supported by the decoder.""" + if self.embed_positions is None: + return self.max_target_positions + return min(self.max_target_positions, + self.embed_positions.max_positions) + + def buffered_future_mask(self, tensor): + dim = tensor.size(0) + if (self._future_mask.size(0) == 0 + or (not self._future_mask.device == tensor.device) + or self._future_mask.size(0) < dim): + self._future_mask = torch.triu( + utils.fill_with_neg_inf(torch.zeros([dim, dim])), 1) + self._future_mask = self._future_mask.to(tensor) + return self._future_mask[:dim, :dim] + + def upgrade_state_dict_named(self, state_dict, name): + """Upgrade a (possibly old) state dict for new versions of fairseq.""" + if isinstance(self.embed_positions, SinusoidalPositionalEmbedding): + weights_key = '{}.embed_positions.weights'.format(name) + if weights_key in state_dict: + del state_dict[weights_key] + state_dict['{}.embed_positions._float_tensor'.format( + name)] = torch.FloatTensor(1) + + if f'{name}.output_projection.weight' not in state_dict: + if self.share_input_output_embed: + embed_out_key = f'{name}.embed_tokens.weight' + else: + embed_out_key = f'{name}.embed_out' + if embed_out_key in state_dict: + state_dict[f'{name}.output_projection.weight'] = state_dict[ + embed_out_key] + if not self.share_input_output_embed: + del state_dict[embed_out_key] + + for i in range(self.num_layers): + # update layer norms + layer_norm_map = { + '0': 'self_attn_layer_norm', + '1': 'encoder_attn_layer_norm', + '2': 'final_layer_norm', + } + for old, new in layer_norm_map.items(): + for m in ('weight', 'bias'): + k = '{}.layers.{}.layer_norms.{}.{}'.format( + name, i, old, m) + if k in state_dict: + state_dict['{}.layers.{}.{}.{}'.format( + name, i, new, m)] = state_dict[k] + del state_dict[k] + + version_key = '{}.version'.format(name) + if utils.item(state_dict.get(version_key, torch.Tensor([1]))[0]) <= 2: + # earlier checkpoints did not normalize after the stack of layers + self.layer_norm = None + self.normalize = False + state_dict[version_key] = torch.Tensor([1]) + + return state_dict + + +def Embedding(num_embeddings, embedding_dim, padding_idx): + m = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx) + nn.init.normal_(m.weight, mean=0, std=embedding_dim**-0.5) + nn.init.constant_(m.weight[padding_idx], 0) + return m + + +def Linear(in_features, out_features, bias=True): + m = nn.Linear(in_features, out_features, bias) + nn.init.xavier_uniform_(m.weight) + if bias: + nn.init.constant_(m.bias, 0.0) + return m + + +def base_architecture(args): + args.encoder_embed_path = getattr(args, 'encoder_embed_path', None) + args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 512) + args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 2048) + args.encoder_layers = getattr(args, 'encoder_layers', 6) + args.encoder_attention_heads = getattr(args, 'encoder_attention_heads', 8) + args.encoder_normalize_before = getattr(args, 'encoder_normalize_before', + False) + args.encoder_learned_pos = getattr(args, 'encoder_learned_pos', False) + args.decoder_embed_path = getattr(args, 'decoder_embed_path', None) + args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', + args.encoder_embed_dim) + args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', + args.encoder_ffn_embed_dim) + args.decoder_layers = getattr(args, 'decoder_layers', 6) + args.decoder_attention_heads = getattr(args, 'decoder_attention_heads', 8) + args.decoder_normalize_before = getattr(args, 'decoder_normalize_before', + False) + args.decoder_learned_pos = getattr(args, 'decoder_learned_pos', False) + args.attention_dropout = getattr(args, 'attention_dropout', 0.0) + args.activation_dropout = getattr(args, 'activation_dropout', 0.0) + args.activation_fn = getattr(args, 'activation_fn', 'relu') + args.dropout = getattr(args, 'dropout', 0.1) + args.adaptive_softmax_cutoff = getattr(args, 'adaptive_softmax_cutoff', + None) + args.adaptive_softmax_dropout = getattr(args, 'adaptive_softmax_dropout', + 0) + args.share_decoder_input_output_embed = getattr( + args, 'share_decoder_input_output_embed', False) + args.share_all_embeddings = getattr(args, 'share_all_embeddings', False) + args.no_token_positional_embeddings = getattr( + args, 'no_token_positional_embeddings', False) + args.adaptive_input = getattr(args, 'adaptive_input', False) + args.no_cross_attention = getattr(args, 'no_cross_attention', False) + args.cross_self_attention = getattr(args, 'cross_self_attention', False) + + args.decoder_output_dim = getattr(args, 'decoder_output_dim', + args.decoder_embed_dim) + args.decoder_input_dim = getattr(args, 'decoder_input_dim', + args.decoder_embed_dim) + + args.no_scale_embedding = getattr(args, 'no_scale_embedding', False) + args.layernorm_embedding = getattr(args, 'layernorm_embedding', False) + args.tie_adaptive_weights = getattr(args, 'tie_adaptive_weights', False) + args.checkpoint_activations = getattr(args, 'checkpoint_activations', + False) + args.offload_activations = getattr(args, 'offload_activations', False) + if args.offload_activations: + args.checkpoint_activations = True + args.encoder_layers_to_keep = getattr(args, 'encoder_layers_to_keep', None) + args.decoder_layers_to_keep = getattr(args, 'decoder_layers_to_keep', None) + args.encoder_layerdrop = getattr(args, 'encoder_layerdrop', 0) + args.decoder_layerdrop = getattr(args, 'decoder_layerdrop', 0) + args.quant_noise_pq = getattr(args, 'quant_noise_pq', 0) + args.quant_noise_pq_block_size = getattr(args, 'quant_noise_pq_block_size', + 8) + args.quant_noise_scalar = getattr(args, 'quant_noise_scalar', 0) + + +def transformer_deep(args): + args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 768) + args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 3072) + args.encoder_attention_heads = getattr(args, 'encoder_attention_heads', 12) + args.encoder_layers = getattr(args, 'encoder_layers', 24) + args.encoder_normalize_before = getattr(args, 'encoder_normalize_before', + True) + args.decoder_normalize_before = getattr(args, 'decoder_normalize_before', + True) + args.decoder_layers = getattr(args, 'decoder_layers', 3) + args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 768) + args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 3072) + args.decoder_attention_heads = getattr(args, 'decoder_attention_heads', 12) + args.attention_dropout = getattr(args, 'attention_dropout', 0.01) + args.activation_dropout = getattr(args, 'activation_dropout', 0.01) + args.dropout = getattr(args, 'dropout', 0.01) + base_architecture(args) diff --git a/modelscope/models/nlp/canmt/canmt_translation.py b/modelscope/models/nlp/canmt/canmt_translation.py new file mode 100644 index 00000000..fe9b933f --- /dev/null +++ b/modelscope/models/nlp/canmt/canmt_translation.py @@ -0,0 +1,78 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import math +import os.path as osp +from typing import Any, Dict, List, Optional, Tuple + +import numpy +import torch +import torch.nn as nn +from torch import Tensor + +from modelscope.metainfo import Models +from modelscope.models.base import TorchModel +from modelscope.models.builder import MODELS +from modelscope.utils.config import Config +from modelscope.utils.constant import ModelFile, Tasks + +__all__ = ['CanmtForTranslation'] + + +@MODELS.register_module( + Tasks.competency_aware_translation, module_name=Models.canmt) +class CanmtForTranslation(TorchModel): + + def __init__(self, model_dir, **args): + """ + CanmtForTranslation implements a Competency-Aware Neural Machine Translaton, + which has both translation and self-estimation abilities. + + For more details, please refer to https://aclanthology.org/2022.emnlp-main.330.pdf + """ + super().__init__(model_dir=model_dir, **args) + self.args = args + cfg_file = osp.join(model_dir, ModelFile.CONFIGURATION) + self.cfg = Config.from_file(cfg_file) + from fairseq.data import Dictionary + self.vocab_src = Dictionary.load(osp.join(model_dir, 'dict.src.txt')) + self.vocab_tgt = Dictionary.load(osp.join(model_dir, 'dict.tgt.txt')) + self.model = self.build_model(model_dir) + self.generator = self.build_generator(self.model, self.vocab_tgt, + self.cfg['decode']) + + def build_model(self, model_dir): + from .canmt_model import CanmtModel + state = self.load_checkpoint( + osp.join(model_dir, ModelFile.TORCH_MODEL_FILE), 'cpu') + cfg = state['cfg'] + model = CanmtModel.build_model(cfg['model'], self) + model.load_state_dict(state['model'], model_cfg=cfg['model']) + return model + + def build_generator(cls, model, vocab_tgt, args): + from .sequence_generator import SequenceGenerator + return SequenceGenerator( + model, + vocab_tgt, + beam_size=args['beam'], + len_penalty=args['lenpen']) + + def load_checkpoint(self, path: str, device: torch.device): + state_dict = torch.load(path, map_location=device) + self.load_state_dict(state_dict, strict=False) + return state_dict + + def forward(self, input: Dict[str, Dict]): + """return the result by the model + + Args: + input (Dict[str, Tensor]): the preprocessed data which contains following: + - src_tokens: tensor with shape (2478,242,24,4), + - src_lengths: tensor with shape (4) + + + Returns: + Dict[str, Tensor]: results which contains following: + - predictions: tokens need to be decode by tokenizer with shape [1377, 4959, 2785, 6392...] + """ + input = {'net_input': input} + return self.generator.generate(input) diff --git a/modelscope/models/nlp/canmt/sequence_generator.py b/modelscope/models/nlp/canmt/sequence_generator.py new file mode 100644 index 00000000..8a5134a4 --- /dev/null +++ b/modelscope/models/nlp/canmt/sequence_generator.py @@ -0,0 +1,850 @@ +# Part of the implementation is borrowed and modified from FAIRSEQ, +# publicly available at https://github.com/facebookresearch/fairseq +# Copyright 2022-2023 The Alibaba MT Team Authors. All rights reserved. +import math +import sys +from typing import Dict, List, Optional + +import numpy +import torch +import torch.nn as nn +from fairseq import search, utils +from fairseq.data import data_utils +from fairseq.models import FairseqIncrementalDecoder +from fairseq.ngram_repeat_block import NGramRepeatBlock +from torch import Tensor + + +def label_smoothed_nll_loss(lprobs, + target, + epsilon, + ignore_index=None, + reduce=True): + if target.dim() == lprobs.dim() - 1: + target = target.unsqueeze(-1) + if target.dtype != torch.int64: + target = target.type(torch.int64) + nll_loss = -lprobs.gather(dim=-1, index=target) + smooth_loss = -lprobs.sum(dim=-1, keepdim=True) + if ignore_index is not None: + pad_mask = target.eq(ignore_index) + nll_loss.masked_fill_(pad_mask, 0.0) + smooth_loss.masked_fill_(pad_mask, 0.0) + else: + nll_loss = nll_loss.squeeze(-1) + smooth_loss = smooth_loss.squeeze(-1) + if reduce: + nll_loss = nll_loss.sum() + smooth_loss = smooth_loss.sum() + eps_i = epsilon / (lprobs.size(-1) - 1) + loss = (1.0 - epsilon - eps_i) * nll_loss + eps_i * smooth_loss + return loss, nll_loss + + +class SequenceGenerator(nn.Module): + + def __init__( + self, + model, + tgt_dict, + beam_size=1, + max_len_a=0, + max_len_b=200, + max_len=200, + min_len=1, + normalize_scores=True, + len_penalty=1.0, + unk_penalty=0.0, + temperature=1.0, + match_source_len=False, + no_repeat_ngram_size=0, + search_strategy=None, + eos=None, + symbols_to_strip_from_output=None, + lm_model=None, + lm_weight=1.0, + recon_force_decoding=True, + trans_force_decoding=False, + ): + """Generates translations of a given source sentence. + + Args: + models (List[~fairseq.models.FairseqModel]): ensemble of models, + currently support fairseq.models.TransformerModel for scripting + beam_size (int, optional): beam width (default: 1) + max_len_a/b (int, optional): generate sequences of maximum length + ax + b, where x is the source length + max_len (int, optional): the maximum length of the generated output + (not including end-of-sentence) + min_len (int, optional): the minimum length of the generated output + (not including end-of-sentence) + normalize_scores (bool, optional): normalize scores by the length + of the output (default: True) + len_penalty (float, optional): length penalty, where <1.0 favors + shorter, >1.0 favors longer sentences (default: 1.0) + unk_penalty (float, optional): unknown word penalty, where <0 + produces more unks, >0 produces fewer (default: 0.0) + temperature (float, optional): temperature, where values + >1.0 produce more uniform samples and values <1.0 produce + sharper samples (default: 1.0) + match_source_len (bool, optional): outputs should match the source + length (default: False) + """ + super().__init__() + self.model = model + self.recon_force_decoding = recon_force_decoding + self.trans_force_decoding = trans_force_decoding + self.tgt_dict = tgt_dict + self.pad = tgt_dict.pad() + self.unk = tgt_dict.unk() + self.eos = tgt_dict.eos() if eos is None else eos + self.symbols_to_strip_from_output = ( + symbols_to_strip_from_output.union({self.eos}) + if symbols_to_strip_from_output is not None else {self.eos}) + self.vocab_size = len(tgt_dict) + self.beam_size = beam_size + # the max beam size is the dictionary size - 1, since we never select pad + self.beam_size = min(beam_size, self.vocab_size - 1) + self.max_len_a = max_len_a + self.max_len_b = max_len_b + self.min_len = min_len + self.max_len = max_len + + self.normalize_scores = normalize_scores + self.len_penalty = len_penalty + self.unk_penalty = unk_penalty + self.temperature = temperature + self.match_source_len = match_source_len + self.eps = 0.1 + + if no_repeat_ngram_size > 0: + self.repeat_ngram_blocker = NGramRepeatBlock(no_repeat_ngram_size) + else: + self.repeat_ngram_blocker = None + + assert temperature > 0, '--temperature must be greater than 0' + + self.search = ( + search.BeamSearch(tgt_dict) + if search_strategy is None else search_strategy) + # We only need to set src_lengths in LengthConstrainedBeamSearch. + # As a module attribute, setting it would break in multithread + # settings when the model is shared. + self.should_set_src_lengths = ( + hasattr(self.search, 'needs_src_lengths') + and self.search.needs_src_lengths) + + self.model.eval() + + self.lm_model = lm_model + self.lm_weight = lm_weight + if self.lm_model is not None: + self.lm_model.eval() + + def cuda(self): + self.model.cuda() + return self + + @torch.no_grad() + def forward( + self, + sample: Dict[str, Dict[str, Tensor]], + prefix_tokens: Optional[Tensor] = None, + bos_token: Optional[int] = None, + ): + """Generate a batch of translations. + + Args: + sample (dict): batch + prefix_tokens (torch.LongTensor, optional): force decoder to begin + with these tokens + bos_token (int, optional): beginning of sentence token + (default: self.eos) + """ + return self._generate(sample, prefix_tokens, bos_token=bos_token) + + # TODO(myleott): unused, deprecate after pytorch-translate migration + def generate_batched_itr(self, + data_itr, + beam_size=None, + cuda=False, + timer=None): + """Iterate over a batched dataset and yield individual translations. + Args: + cuda (bool, optional): use GPU for generation + timer (StopwatchMeter, optional): time generations + """ + for sample in data_itr: + s = utils.move_to_cuda(sample) if cuda else sample + if 'net_input' not in s: + continue + input = s['net_input'] + # model.forward normally channels prev_output_tokens into the decoder + # separately, but SequenceGenerator directly calls model.encoder + encoder_input = { + k: v + for k, v in input.items() if k != 'prev_output_tokens' + } + if timer is not None: + timer.start() + with torch.no_grad(): + hypos = self.generate(encoder_input) + if timer is not None: + timer.stop(sum(len(h[0]['tokens']) for h in hypos)) + for i, id in enumerate(s['id'].data): + # remove padding + src = utils.strip_pad(input['src_tokens'].data[i, :], self.pad) + ref = ( + utils.strip_pad(s['target'].data[i, :], self.pad) + if s['target'] is not None else None) + yield id, src, ref, hypos[i] + + @torch.no_grad() + def generate(self, sample: Dict[str, Dict[str, Tensor]], + **kwargs) -> List[List[Dict[str, Tensor]]]: + """Generate translations. Match the api of other fairseq generators. + + Args: + models (List[~fairseq.models.FairseqModel]): ensemble of models + sample (dict): batch + prefix_tokens (torch.LongTensor, optional): force decoder to begin + with these tokens + constraints (torch.LongTensor, optional): force decoder to include + the list of constraints + bos_token (int, optional): beginning of sentence token + (default: self.eos) + """ + from torch import tensor + finalized = self._generate(sample, **kwargs) + tokens_list = [] + decoder_list = [] + for i in range(len(finalized)): + sent = finalized[i][0] + tokens = sent['tokens'] + tokens_list.append(tokens) + decoder_out = sent['decoder_out'] + decoder_list.append(decoder_out) + # padding tokens + size = max(v.size(0) for v in tokens_list) + batch_size = len(tokens_list) + + for i in range(len(tokens_list)): + tokens_list[i] = tokens_list[i].roll(1, 0) + decoder_list[i] = decoder_list[i].roll(1, 0) + res = tokens_list[0].new(batch_size, size).fill_(self.pad) + + def copy_tensor(src, dst): + assert dst.numel() == src.numel() + dst.copy_(src) + + for i, v in enumerate(tokens_list): + copy_tensor(v, res[i][:len(v)] if True else res[i][:len(v)]) + tgt_tokens = res + decoder_padding_mask = tgt_tokens.eq(self.pad) + + # generate for src reconstruction + decoder_out_re = self.model.decoder( + tgt_tokens, + encoder_out=None, + full_context_alignment=True, + ) + decoder_out_re = decoder_out_re[1]['last_layer'] + decoder_outs = dict() + decoder_outs['encoder_out'] = [decoder_out_re] + decoder_outs['encoder_padding_mask'] = [decoder_padding_mask] + decoder_outs['src_tokens'] = [tgt_tokens] + decoder_outs['encoder_embedding'] = [] + decoder_outs['encoder_states'] = [] + decoder_outs['src_lengths'] = [] + + scores = self._forward_src(decoder_outs, sample) + return finalized, scores + + def _forward( + self, + sample: Dict[str, Dict[str, Tensor]], + ): + net_input = sample['net_input'] + src_tokens = net_input['src_tokens'] + prev_output_tokens = net_input['prev_output_tokens'] + encoder_outs = self.model.forward_encoder(net_input) + final_encoder_out = encoder_outs['encoder_out'][0].transpose(0, 1) + final_encoder_embedding = encoder_outs['encoder_embedding'][0] + self.model.has_incremental = False + lprobs, avg_attn_scores, decoder_outs = self.model.forward_decoder( + prev_output_tokens, encoder_outs, incremental_states=None) + decoder_outs = decoder_outs.transpose(0, 1) + self.model.has_incremental = True + batch_size = src_tokens.size(0) + # list of completed sentences + finalized = torch.jit.annotate( + List[List[Dict[str, Tensor]]], + [ + torch.jit.annotate(List[Dict[str, Tensor]], []) + for i in range(batch_size) + ], + ) # contains lists of dictionaries of infomation about the hypothesis being finalized at each step + for i in range(batch_size): + eos_idx = numpy.where(sample['target'][i].cpu() == self.eos)[0][0] + finalized[i].append({ + 'tokens': + sample['target'][i][:eos_idx + 1], + 'decoder_out': + decoder_outs[i][:eos_idx + 1], + 'final_encoder_embedding': + final_encoder_embedding[i], + 'final_encoder_out': + final_encoder_out[i], + }) + return finalized + + def _forward_src( + self, + encoder_outs, + sample: Dict[str, Dict[str, Tensor]], + ): + net_input = sample['net_input'] + src_tokens = net_input['src_tokens'] + src_lengths = net_input['src_lengths'] + prev_output_tokens = net_input['prev_src_tokens'] + self.model.has_incremental = False + lprobs, avg_attn_scores, _, decoder_outs = self.model.forward_decoder_src( + prev_output_tokens, encoder_outs, incremental_states=None) + logits = decoder_outs[0] + lprobs = utils.log_softmax(logits, dim=-1) + sources = net_input['sources'].view(-1) + lprobs = lprobs.view(-1, lprobs.size(-1)) + + scores, _ = label_smoothed_nll_loss( + lprobs, + sources, + epsilon=self.eps, + ignore_index=self.pad, + reduce=False) + batch_size = src_tokens.shape[0] + scores = scores.reshape(batch_size, -1) + scores = torch.cat((scores.sum(axis=-1, keepdim=True) + / src_lengths.reshape(batch_size, 1), scores), + axis=-1) + return scores + + def _generate( + self, + sample: Dict[str, Dict[str, Tensor]], + prefix_tokens: Optional[Tensor] = None, + constraints: Optional[Tensor] = None, + bos_token: Optional[int] = None, + ): + incremental_states = torch.jit.annotate( + Dict[str, Dict[str, Optional[Tensor]]], + torch.jit.annotate(Dict[str, Dict[str, Optional[Tensor]]], {}), + ) + + net_input = sample['net_input'] + if 'src_tokens' in net_input: + src_tokens = net_input['src_tokens'] + # length of the source text being the character length except EndOfSentence and pad + src_lengths = ((src_tokens.ne(self.eos) + & src_tokens.ne(self.pad)).long().sum(dim=1)) + else: + raise Exception( + 'expected src_tokens or source in net input. input keys: ' + + str(net_input.keys())) + + # bsz: total number of sentences in beam + # Note that src_tokens may have more than 2 dimensions (i.e. audio features) + bsz, src_len = src_tokens.size()[:2] + beam_size = self.beam_size + + max_len: int = -1 + if self.match_source_len: + max_len = src_lengths.max().item() + else: + max_len = min( + int(self.max_len_a * src_len + self.max_len_b), + self.max_len - 1, + ) + assert ( + self.min_len <= max_len + ), 'min_len cannot be larger than max_len, please adjust these!' + # compute the encoder output for each beam + encoder_outs = self.model.forward_encoder(net_input) + + final_encoder_out = encoder_outs['encoder_out'][0].transpose(0, 1) + final_encoder_embedding = encoder_outs['encoder_embedding'][0] + + # placeholder of indices for bsz * beam_size to hold tokens and accumulative scores + new_order = torch.arange(bsz).view(-1, 1).repeat(1, beam_size).view(-1) + new_order = new_order.to(src_tokens.device).long() + encoder_outs = self.model.reorder_encoder_out(encoder_outs, new_order) + + # ensure encoder_outs is a List. + assert encoder_outs is not None + + # initialize buffers + scores = (torch.zeros(bsz * beam_size, + max_len + 1).to(src_tokens).float() + ) # +1 for eos; pad is never chosen for scoring + tokens = (torch.zeros(bsz * beam_size, + max_len + 2).to(src_tokens).long().fill_( + self.pad)) # +2 for eos and pad + tokens[:, 0] = self.eos if bos_token is None else bos_token + attn: Optional[Tensor] = None + + cands_to_ignore = (torch.zeros(bsz, beam_size).to(src_tokens).eq(-1) + ) # forward and backward-compatible False mask + + # list of completed sentences + finalized = torch.jit.annotate( + List[List[Dict[str, Tensor]]], + [ + torch.jit.annotate(List[Dict[str, Tensor]], []) + for i in range(bsz) + ], + ) # contains lists of dictionaries of infomation about the hypothesis being finalized at each step + + # a boolean array indicating if the sentence at the index is finished or not + finished = [False for i in range(bsz)] + num_remaining_sent = bsz # number of sentences remaining + + # number of candidate hypos per step + cand_size = 2 * beam_size # 2 x beam size in case half are EOS + + # offset arrays for converting between different indexing schemes + bbsz_offsets = ((torch.arange(0, bsz) + * beam_size).unsqueeze(1).type_as(tokens).to( + src_tokens.device)) + cand_offsets = torch.arange(0, cand_size).type_as(tokens).to( + src_tokens.device) + + reorder_state: Optional[Tensor] = None + batch_idxs: Optional[Tensor] = None + + original_batch_idxs: Optional[Tensor] = None + if 'id' in sample and isinstance(sample['id'], Tensor): + original_batch_idxs = sample['id'] + else: + original_batch_idxs = torch.arange(0, bsz).type_as(tokens) + + for step in range(max_len + 1): # one extra step for EOS marker + # reorder decoder internal states based on the prev choice of beams + if reorder_state is not None: + if batch_idxs is not None: + # update beam indices to take into account removed sentences + corr = batch_idxs - torch.arange( + batch_idxs.numel()).type_as(batch_idxs) + reorder_state.view(-1, beam_size).add_( + corr.unsqueeze(-1) * beam_size) + original_batch_idxs = original_batch_idxs[batch_idxs] + + self.model.reorder_incremental_state(incremental_states, + reorder_state) + encoder_outs = self.model.reorder_encoder_out( + encoder_outs, reorder_state) + + lprobs, avg_attn_scores, decoder_out_word = self.model.forward_decoder( + tokens[:, :step + 1], + encoder_outs, + incremental_states, + self.temperature, + ) + # (length, batch*beam, hidden_size) - >(batch*beam, length, hidden_size) + decoder_out_word = decoder_out_word.transpose(0, 1) + if step == 0: + decoder_out_tensor = decoder_out_word + else: + decoder_out_tensor = torch.cat( + [decoder_out_tensor, decoder_out_word], dim=1) + + lprobs[lprobs != lprobs] = torch.tensor(-math.inf).to(lprobs) + + lprobs[:, self.pad] = -math.inf # never select pad + lprobs[:, self.unk] -= self.unk_penalty # apply unk penalty + + # handle max length constraint + if step >= max_len: + lprobs[:, :self.eos] = -math.inf + lprobs[:, self.eos + 1:] = -math.inf + + # handle prefix tokens (possibly with different lengths) + if (prefix_tokens is not None and step < prefix_tokens.size(1) + and step < max_len): + + lprobs, tokens, scores = self._prefix_tokens( + step, lprobs, scores, tokens, prefix_tokens, beam_size) + elif step < self.min_len: + # minimum length constraint (does not apply if using prefix_tokens) + lprobs[:, self.eos] = -math.inf + + # Record attention scores, only support avg_attn_scores is a Tensor + if avg_attn_scores is not None: + if attn is None: + attn = torch.empty(bsz * beam_size, + avg_attn_scores.size(1), + max_len + 2).to(scores) + attn[:, :, step + 1].copy_(avg_attn_scores) + + scores = scores.type_as(lprobs) + eos_bbsz_idx = torch.empty(0).to( + tokens + ) # indices of hypothesis ending with eos (finished sentences) + eos_scores = torch.empty(0).to( + scores + ) # scores of hypothesis ending with eos (finished sentences) + + if self.should_set_src_lengths: + self.search.set_src_lengths(src_lengths) + + if self.repeat_ngram_blocker is not None: + lprobs = self.repeat_ngram_blocker(tokens, lprobs, bsz, + beam_size, step) + + # Shape: (batch, cand_size) + cand_scores, cand_indices, cand_beams = self.search.step( + step, + lprobs.view(bsz, -1, self.vocab_size), + scores.view(bsz, beam_size, -1)[:, :, :step], + tokens[:, :step + 1], + original_batch_idxs, + ) + + # cand_bbsz_idx contains beam indices for the top candidate + # hypotheses, with a range of values: [0, bsz*beam_size), + # and dimensions: [bsz, cand_size] + cand_bbsz_idx = cand_beams.add(bbsz_offsets) + + # finalize hypotheses that end in eos + # Shape of eos_mask: (batch size, beam size) + eos_mask = cand_indices.eq(self.eos) & cand_scores.ne(-math.inf) + eos_mask[:, :beam_size][cands_to_ignore] = torch.tensor(0).to( + eos_mask) + + # only consider eos when it's among the top beam_size indices + # Now we know what beam item(s) to finish + # Shape: 1d list of absolute-numbered + eos_bbsz_idx = torch.masked_select( + cand_bbsz_idx[:, :beam_size], mask=eos_mask[:, :beam_size]) + + finalized_sents: List[int] = [] + if eos_bbsz_idx.numel() > 0: + eos_scores = torch.masked_select( + cand_scores[:, :beam_size], mask=eos_mask[:, :beam_size]) + + finalized_sents = self.finalize_hypos( + step, + eos_bbsz_idx, + eos_scores, + tokens, + scores, + finalized, + finished, + beam_size, + attn, + src_lengths, + max_len, + decoder_out_tensor, + ) + num_remaining_sent -= len(finalized_sents) + + assert num_remaining_sent >= 0 + if num_remaining_sent == 0: + break + if self.search.stop_on_max_len and step >= max_len: + break + assert step < max_len, f'{step} < {max_len}' + + # Remove finalized sentences (ones for which {beam_size} + # finished hypotheses have been generated) from the batch. + if len(finalized_sents) > 0: + new_bsz = bsz - len(finalized_sents) + + # construct batch_idxs which holds indices of batches to keep for the next pass + batch_mask = torch.ones( + bsz, dtype=torch.bool, device=cand_indices.device) + batch_mask[finalized_sents] = False + # TODO replace `nonzero(as_tuple=False)` after TorchScript supports it + batch_idxs = torch.arange( + bsz, device=cand_indices.device).masked_select(batch_mask) + + # Choose the subset of the hypothesized constraints that will continue + self.search.prune_sentences(batch_idxs) + + eos_mask = eos_mask[batch_idxs] + cand_beams = cand_beams[batch_idxs] + bbsz_offsets.resize_(new_bsz, 1) + cand_bbsz_idx = cand_beams.add(bbsz_offsets) + cand_scores = cand_scores[batch_idxs] + cand_indices = cand_indices[batch_idxs] + + if prefix_tokens is not None: + prefix_tokens = prefix_tokens[batch_idxs] + src_lengths = src_lengths[batch_idxs] + cands_to_ignore = cands_to_ignore[batch_idxs] + + scores = scores.view(bsz, -1)[batch_idxs].view( + new_bsz * beam_size, -1) + tokens = tokens.view(bsz, -1)[batch_idxs].view( + new_bsz * beam_size, -1) + decoder_out_tensor = decoder_out_tensor.contiguous().view( + bsz, -1)[batch_idxs].view(new_bsz * beam_size, + decoder_out_tensor.size(1), -1) + if attn is not None: + attn = attn.view(bsz, -1)[batch_idxs].view( + new_bsz * beam_size, attn.size(1), -1) + bsz = new_bsz + else: + batch_idxs = None + + # Set active_mask so that values > cand_size indicate eos hypos + # and values < cand_size indicate candidate active hypos. + # After, the min values per row are the top candidate active hypos + + eos_mask[:, :beam_size] = ~((~cands_to_ignore) + & (~eos_mask[:, :beam_size])) + active_mask = torch.add( + eos_mask.type_as(cand_offsets) * cand_size, + cand_offsets[:eos_mask.size(1)], + ) + + # get the top beam_size active hypotheses, which are just + # the hypos with the smallest values in active_mask. + # {active_hypos} indicates which {beam_size} hypotheses + # from the list of {2 * beam_size} candidates were + # selected. Shapes: (batch size, beam size) + new_cands_to_ignore, active_hypos = torch.topk( + active_mask, k=beam_size, dim=1, largest=False) + + # update cands_to_ignore to ignore any finalized hypos. + cands_to_ignore = new_cands_to_ignore.ge(cand_size)[:, :beam_size] + # Make sure there is at least one active item for each sentence in the batch. + assert (~cands_to_ignore).any(dim=1).all() + + # update cands_to_ignore to ignore any finalized hypos + + # {active_bbsz_idx} denotes which beam number is continued for each new hypothesis (a beam + # can be selected more than once). + active_bbsz_idx = torch.gather( + cand_bbsz_idx, dim=1, index=active_hypos) + active_scores = torch.gather( + cand_scores, dim=1, index=active_hypos) + + active_bbsz_idx = active_bbsz_idx.view(-1) + active_scores = active_scores.view(-1) + + # copy tokens and scores for active hypotheses + + # Set the tokens for each beam (can select the same row more than once) + tokens[:, :step + 1] = torch.index_select( + tokens[:, :step + 1], dim=0, index=active_bbsz_idx) + # Select the next token for each of them + tokens.view(bsz, beam_size, -1)[:, :, step + 1] = torch.gather( + cand_indices, dim=1, index=active_hypos) + if step > 0: + scores[:, :step] = torch.index_select( + scores[:, :step], dim=0, index=active_bbsz_idx) + scores.view(bsz, beam_size, -1)[:, :, step] = torch.gather( + cand_scores, dim=1, index=active_hypos) + + # Update constraints based on which candidates were selected for the next beam + self.search.update_constraints(active_hypos) + + # copy attention for active hypotheses + if attn is not None: + attn[:, :, :step + 2] = torch.index_select( + attn[:, :, :step + 2], dim=0, index=active_bbsz_idx) + + # reorder incremental state in decoder + reorder_state = active_bbsz_idx + + # sort by score descending + for sent in range(len(finalized)): + scores = torch.tensor( + [float(elem['score'].item()) for elem in finalized[sent]]) + _, sorted_scores_indices = torch.sort(scores, descending=True) + finalized[sent] = [ + finalized[sent][ssi] for ssi in sorted_scores_indices + ] + finalized[sent] = torch.jit.annotate(List[Dict[str, Tensor]], + finalized[sent]) + finalized[sent][0][ + 'final_encoder_embedding'] = final_encoder_embedding[sent] + finalized[sent][0]['final_encoder_out'] = final_encoder_out[sent] + return finalized + + def _prefix_tokens(self, step: int, lprobs, scores, tokens, prefix_tokens, + beam_size: int): + """Handle prefix tokens""" + prefix_toks = prefix_tokens[:, step].unsqueeze(-1).repeat( + 1, beam_size).view(-1) + prefix_lprobs = lprobs.gather(-1, prefix_toks.unsqueeze(-1)) + prefix_mask = prefix_toks.ne(self.pad) + lprobs[prefix_mask] = torch.tensor(-math.inf).to(lprobs) + lprobs[prefix_mask] = lprobs[prefix_mask].scatter( + -1, prefix_toks[prefix_mask].unsqueeze(-1), + prefix_lprobs[prefix_mask]) + # if prefix includes eos, then we should make sure tokens and + # scores are the same across all beams + eos_mask = prefix_toks.eq(self.eos) + if eos_mask.any(): + # validate that the first beam matches the prefix + first_beam = tokens[eos_mask].view(-1, beam_size, + tokens.size(-1))[:, 0, + 1:step + 1] + eos_mask_batch_dim = eos_mask.view(-1, beam_size)[:, 0] + target_prefix = prefix_tokens[eos_mask_batch_dim][:, :step] + assert (first_beam == target_prefix).all() + + # copy tokens, scores and lprobs from the first beam to all beams + tokens = self.replicate_first_beam(tokens, eos_mask_batch_dim, + beam_size) + scores = self.replicate_first_beam(scores, eos_mask_batch_dim, + beam_size) + lprobs = self.replicate_first_beam(lprobs, eos_mask_batch_dim, + beam_size) + return lprobs, tokens, scores + + def replicate_first_beam(self, tensor, mask, beam_size: int): + tensor = tensor.view(-1, beam_size, tensor.size(-1)) + tensor[mask] = tensor[mask][:, :1, :] + return tensor.view(-1, tensor.size(-1)) + + def finalize_hypos( + self, + step: int, + bbsz_idx, + eos_scores, + tokens, + scores, + finalized: List[List[Dict[str, Tensor]]], + finished: List[bool], + beam_size: int, + attn: Optional[Tensor], + src_lengths, + max_len: int, + decoder_out=None, + ): + """Finalize hypothesis, store finalized information in `finalized`, and change `finished` accordingly. + A sentence is finalized when {beam_size} finished items have been collected for it. + + Returns number of sentences (not beam items) being finalized. + These will be removed from the batch and not processed further. + Args: + bbsz_idx (Tensor): + """ + assert bbsz_idx.numel() == eos_scores.numel() + + # clone relevant token and attention tensors. + # tokens is (batch * beam, max_len). So the index_select + # gets the newly EOS rows, then selects cols 1..{step + 2} + if decoder_out is not None: + decoder_out_clone = decoder_out.index_select(0, bbsz_idx) + + tokens_clone = tokens.index_select(0, bbsz_idx)[:, 1:step + 2] + # skip the first index, which is EOS + + tokens_clone[:, step] = self.eos + attn_clone = ( + attn.index_select(0, bbsz_idx)[:, :, 1:step + + 2] if attn is not None else None) + + # compute scores per token position + pos_scores = scores.index_select(0, bbsz_idx)[:, :step + 1] + pos_scores[:, step] = eos_scores + # convert from cumulative to per-position scores + pos_scores[:, 1:] = pos_scores[:, 1:] - pos_scores[:, :-1] + + # normalize sentence-level scores + if self.normalize_scores: + eos_scores /= (step + 1)**self.len_penalty + + # cum_unfin records which sentences in the batch are finished. + # It helps match indexing between (a) the original sentences + # in the batch and (b) the current, possibly-reduced set of + # sentences. + cum_unfin: List[int] = [] + prev = 0 + for f in finished: + if f: + prev += 1 + else: + cum_unfin.append(prev) + + # The keys here are of the form "{sent}_{unfin_idx}", where + # "unfin_idx" is the index in the current (possibly reduced) + # list of sentences, and "sent" is the index in the original, + # unreduced batch + # set() is not supported in script export + sents_seen: Dict[str, Optional[Tensor]] = {} + + # For every finished beam item + for i in range(bbsz_idx.size()[0]): + idx = bbsz_idx[i] + score = eos_scores[i] + # sentence index in the current (possibly reduced) batch + unfin_idx = idx // beam_size + # sentence index in the original (unreduced) batch + sent = unfin_idx + cum_unfin[unfin_idx] + # Cannot create dict for key type '(int, int)' in torchscript. + # The workaround is to cast int to string + seen = str(sent.item()) + '_' + str(unfin_idx.item()) + if seen not in sents_seen: + sents_seen[seen] = None + + if self.match_source_len and step > src_lengths[unfin_idx]: + score = torch.tensor(-math.inf).to(score) + + # An input sentence (among those in a batch) is finished when + # beam_size hypotheses have been collected for it + if len(finalized[sent]) < beam_size: + if attn_clone is not None: + # remove padding tokens from attn scores + hypo_attn = attn_clone[i] + else: + hypo_attn = torch.empty(0) + + finalized[sent].append({ + 'tokens': + tokens_clone[i], + 'score': + score, + 'attention': + hypo_attn, # src_len x tgt_len + 'alignment': + torch.empty(0), + 'positional_scores': + pos_scores[i], + 'decoder_out': + decoder_out_clone[i] if decoder_out is not None else [], + }) + + newly_finished: List[int] = [] + + for seen in sents_seen.keys(): + # check termination conditions for this sentence + sent: int = int(float(seen.split('_')[0])) + unfin_idx: int = int(float(seen.split('_')[1])) + + if not finished[sent] and self.is_finished( + step, unfin_idx, max_len, len(finalized[sent]), beam_size): + finished[sent] = True + newly_finished.append(unfin_idx) + + return newly_finished + + def is_finished( + self, + step: int, + unfin_idx: int, + max_len: int, + finalized_sent_len: int, + beam_size: int, + ): + """ + Check whether decoding for a sentence is finished, which + occurs when the list of finalized sentences has reached the + beam size, or when we reach the maximum length. + """ + assert finalized_sent_len <= beam_size + if finalized_sent_len == beam_size or step == max_len: + return True + return False diff --git a/modelscope/pipeline_inputs.py b/modelscope/pipeline_inputs.py index 48afb2b5..a5db3b39 100644 --- a/modelscope/pipeline_inputs.py +++ b/modelscope/pipeline_inputs.py @@ -195,6 +195,8 @@ TASK_INPUTS = { InputType.TEXT, Tasks.translation: InputType.TEXT, + Tasks.competency_aware_translation: + InputType.TEXT, Tasks.word_segmentation: [InputType.TEXT, { 'text': InputType.TEXT, }], diff --git a/modelscope/pipelines/nlp/__init__.py b/modelscope/pipelines/nlp/__init__.py index cc8487fe..d5898f92 100644 --- a/modelscope/pipelines/nlp/__init__.py +++ b/modelscope/pipelines/nlp/__init__.py @@ -30,6 +30,7 @@ if TYPE_CHECKING: from .fid_dialogue_pipeline import FidDialoguePipeline from .token_classification_pipeline import TokenClassificationPipeline from .translation_pipeline import TranslationPipeline + from .canmt_translation_pipeline import CanmtTranslationPipeline from .word_segmentation_pipeline import WordSegmentationPipeline, WordSegmentationThaiPipeline from .zero_shot_classification_pipeline import ZeroShotClassificationPipeline from .mglm_text_summarization_pipeline import MGLMTextSummarizationPipeline @@ -78,6 +79,7 @@ else: 'fid_dialogue_pipeline': ['FidDialoguePipeline'], 'token_classification_pipeline': ['TokenClassificationPipeline'], 'translation_pipeline': ['TranslationPipeline'], + 'canmt_translation_pipeline': ['CanmtTranslationPipeline'], 'translation_quality_estimation_pipeline': ['TranslationQualityEstimationPipeline'], 'word_segmentation_pipeline': diff --git a/modelscope/pipelines/nlp/canmt_translation_pipeline.py b/modelscope/pipelines/nlp/canmt_translation_pipeline.py new file mode 100644 index 00000000..a31399f9 --- /dev/null +++ b/modelscope/pipelines/nlp/canmt_translation_pipeline.py @@ -0,0 +1,91 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import os.path as osp +from typing import Any, Dict, Optional, Union + +import torch +from sacremoses import MosesDetokenizer + +from modelscope.metainfo import Pipelines +from modelscope.models import Model +from modelscope.models.nlp import CanmtForTranslation +from modelscope.outputs import OutputKeys +from modelscope.pipelines.base import Pipeline, Tensor +from modelscope.pipelines.builder import PIPELINES +from modelscope.preprocessors import CanmtTranslationPreprocessor, Preprocessor +from modelscope.utils.constant import ModelFile, Tasks + +__all__ = ['CanmtTranslationPipeline'] + + +@PIPELINES.register_module( + Tasks.competency_aware_translation, + module_name=Pipelines.canmt_translation) +class CanmtTranslationPipeline(Pipeline): + + def __init__(self, + model: Union[Model, str], + preprocessor: Optional[Preprocessor] = None, + config_file: str = None, + device: str = 'gpu', + auto_collate=True, + **kwargs): + """Use `model` and `preprocessor` to create a canmt translation pipeline for prediction. + + Args: + model (str or Model): Supply either a local model dir which supported the canmt translation task, + or a model id from the model hub, or a torch model instance. + preprocessor (Preprocessor): An optional preprocessor instance, please make sure the preprocessor fits for + the model if supplied. + kwargs (dict, `optional`): + Extra kwargs passed into the preprocessor's constructor. + + Examples: + >>> from modelscope.pipelines import pipeline + >>> pipeline_ins = pipeline(task='competency_aware_translation', + >>> model='damo/nlp_canmt_translation_zh2en_large') + >>> sentence1 = '世界是丰富多彩的。' + >>> print(pipeline_ins(sentence1)) + >>> # Or use the list input: + >>> print(pipeline_ins([sentence1]) + + To view other examples plese check tests/pipelines/test_canmt_translation.py. + """ + super().__init__( + model=model, + preprocessor=preprocessor, + config_file=config_file, + device=device, + auto_collate=auto_collate) + assert isinstance(self.model, Model), \ + f'please check whether model config exists in {ModelFile.CONFIGURATION}' + + if self.preprocessor is None: + self.preprocessor = CanmtTranslationPreprocessor( + self.model.model_dir, + kwargs) if preprocessor is None else preprocessor + self.vocab_tgt = self.preprocessor.vocab_tgt + self.detokenizer = MosesDetokenizer(lang=self.preprocessor.tgt_lang) + + def forward(self, inputs: Dict[str, Any], + **forward_params) -> Dict[str, Any]: + with torch.no_grad(): + return super().forward(inputs, **forward_params) + + def postprocess(self, inputs: Dict[str, Tensor], + **postprocess_params) -> Dict[str, str]: + batch_size = len(inputs[0]) + hypos = [] + scores = [] + for i in range(batch_size): + hypo_tensor = inputs[0][i][0]['tokens'] + score = inputs[1][i][0].cpu().tolist() + hypo_sent = self.vocab_tgt.string( + hypo_tensor, + '@@ ', + extra_symbols_to_ignore={self.vocab_tgt.pad()}) + hypo_sent = self.detokenizer.detokenize(hypo_sent.split()) + hypos.append(hypo_sent) + scores.append(score) + + return {OutputKeys.TRANSLATION: hypos, OutputKeys.SCORE: scores} diff --git a/modelscope/preprocessors/__init__.py b/modelscope/preprocessors/__init__.py index ee342865..06833b82 100644 --- a/modelscope/preprocessors/__init__.py +++ b/modelscope/preprocessors/__init__.py @@ -40,7 +40,7 @@ if TYPE_CHECKING: DialogStateTrackingPreprocessor, ConversationalTextToSqlPreprocessor, TableQuestionAnsweringPreprocessor, NERPreprocessorViet, NERPreprocessorThai, WordSegmentationPreprocessorThai, - TranslationEvaluationPreprocessor, + TranslationEvaluationPreprocessor, CanmtTranslationPreprocessor, DialogueClassificationUsePreprocessor, SiameseUiePreprocessor, DocumentGroundedDialogGeneratePreprocessor, DocumentGroundedDialogRetrievalPreprocessor, @@ -94,6 +94,7 @@ else: 'ConversationalTextToSqlPreprocessor', 'TableQuestionAnsweringPreprocessor', 'TranslationEvaluationPreprocessor', + 'CanmtTranslationPreprocessor', 'DialogueClassificationUsePreprocessor', 'SiameseUiePreprocessor', 'DialogueClassificationUsePreprocessor', 'DocumentGroundedDialogGeneratePreprocessor', diff --git a/modelscope/preprocessors/base.py b/modelscope/preprocessors/base.py index 951e5c3e..127fc23e 100644 --- a/modelscope/preprocessors/base.py +++ b/modelscope/preprocessors/base.py @@ -15,6 +15,8 @@ logger = get_logger() PREPROCESSOR_MAP = { # nlp + (Models.canmt, Tasks.competency_aware_translation): + Preprocessors.canmt_translation, # bart (Models.bart, Tasks.text_error_correction): Preprocessors.text_error_correction, diff --git a/modelscope/preprocessors/nlp/__init__.py b/modelscope/preprocessors/nlp/__init__.py index f0660374..afd8541c 100644 --- a/modelscope/preprocessors/nlp/__init__.py +++ b/modelscope/preprocessors/nlp/__init__.py @@ -30,6 +30,7 @@ if TYPE_CHECKING: from .space_T_cn import TableQuestionAnsweringPreprocessor from .mglm_summarization_preprocessor import MGLMSummarizationPreprocessor from .translation_evaluation_preprocessor import TranslationEvaluationPreprocessor + from .canmt_translation import CanmtTranslationPreprocessor from .dialog_classification_use_preprocessor import DialogueClassificationUsePreprocessor from .siamese_uie_preprocessor import SiameseUiePreprocessor from .document_grounded_dialog_generate_preprocessor import DocumentGroundedDialogGeneratePreprocessor @@ -90,6 +91,9 @@ else: 'space_T_cn': ['TableQuestionAnsweringPreprocessor'], 'translation_evaluation_preprocessor': ['TranslationEvaluationPreprocessor'], + 'canmt_translation': [ + 'CanmtTranslationPreprocessor', + ], 'dialog_classification_use_preprocessor': ['DialogueClassificationUsePreprocessor'], 'siamese_uie_preprocessor': ['SiameseUiePreprocessor'], diff --git a/modelscope/preprocessors/nlp/canmt_translation.py b/modelscope/preprocessors/nlp/canmt_translation.py new file mode 100644 index 00000000..760afc26 --- /dev/null +++ b/modelscope/preprocessors/nlp/canmt_translation.py @@ -0,0 +1,109 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import os.path as osp +from typing import Any, Dict + +import jieba +import torch +from sacremoses import MosesDetokenizer, MosesPunctNormalizer, MosesTokenizer +from subword_nmt import apply_bpe + +from modelscope.metainfo import Preprocessors +from modelscope.preprocessors.base import Preprocessor +from modelscope.preprocessors.builder import PREPROCESSORS +from modelscope.utils.config import Config +from modelscope.utils.constant import Fields, ModelFile +from .text_clean import TextClean + + +@PREPROCESSORS.register_module( + Fields.nlp, module_name=Preprocessors.canmt_translation) +class CanmtTranslationPreprocessor(Preprocessor): + """The preprocessor used in text correction task. + """ + + def __init__(self, + model_dir: str, + max_length: int = None, + *args, + **kwargs): + from fairseq.data import Dictionary + """preprocess the data via the vocab file from the `model_dir` path + + Args: + model_dir (str): model path + """ + super().__init__(*args, **kwargs) + self.cfg = Config.from_file( + osp.join(model_dir, ModelFile.CONFIGURATION)) + self.vocab_src = Dictionary.load(osp.join(model_dir, 'dict.src.txt')) + self.vocab_tgt = Dictionary.load(osp.join(model_dir, 'dict.tgt.txt')) + self.padding_value = self.vocab_src.pad() + self.max_length = max_length + 1 if max_length is not None else 129 # 1 is eos token + + self.src_lang = self.cfg['preprocessor']['src_lang'] + self.tgt_lang = self.cfg['preprocessor']['tgt_lang'] + self.tc = TextClean() + + if self.src_lang == 'zh': + self.tok = jieba + else: + self.punct_normalizer = MosesPunctNormalizer(lang=self.src_lang) + self.tok = MosesTokenizer(lang=self.src_lang) + + self.src_bpe_path = osp.join( + model_dir, self.cfg['preprocessor']['src_bpe']['file']) + self.bpe = apply_bpe.BPE(open(self.src_bpe_path)) + + def __call__(self, input: str) -> Dict[str, Any]: + """process the raw input data + + Args: + data (str): a sentence + Example: + '随着中国经济突飞猛近，建造工业与日俱增' + Returns: + Dict[str, Any]: the preprocessed data + Example: + {'net_input': + {'src_tokens':tensor([1,2,3,4]), + 'src_lengths': tensor([4])} + } + """ + if self.src_lang == 'zh': + input = self.tc.clean(input) + input_tok = self.tok.cut(input) + input_tok = ' '.join(list(input_tok)) + else: + input = [self._punct_normalizer.normalize(item) for item in input] + input_tok = [ + self.tok.tokenize( + item, return_str=True, aggressive_dash_splits=True) + for item in input + ] + + input_bpe = self.bpe.process_line(input_tok).strip().split() + text = ' '.join([x for x in input_bpe]) + + inputs = self.vocab_src.encode_line( + text, append_eos=True, add_if_not_exist=False) + prev_inputs = torch.roll(inputs, shifts=1) + lengths = inputs.size()[0] + max_len = min(self.max_length, lengths) + + padding = torch.tensor( + [self.padding_value] * # noqa: W504 + (max_len - lengths), + dtype=inputs.dtype) + sources = torch.unsqueeze(torch.cat([inputs, padding]), dim=0) + inputs = torch.unsqueeze(torch.cat([padding, inputs]), dim=0) + prev_inputs = torch.unsqueeze(torch.cat([prev_inputs, padding]), dim=0) + lengths = torch.tensor([lengths]) + out = { + 'src_tokens': inputs, + 'src_lengths': lengths, + 'prev_src_tokens': prev_inputs, + 'sources': sources + } + + return out diff --git a/modelscope/preprocessors/nlp/text_clean.py b/modelscope/preprocessors/nlp/text_clean.py new file mode 100644 index 00000000..a369a141 --- /dev/null +++ b/modelscope/preprocessors/nlp/text_clean.py @@ -0,0 +1,70 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import codecs +import re +import sys + + +class TextClean(object): + + def __init__(self): + spu = [ + 0xA0, 0x1680, 0x202f, 0x205F, 0x3000, 0xFEFF, 8203, 8206, 8207, + 8298, 8300, 65279 + ] + spu.extend(range(0xE000, 0xF8FF + 1)) + spu.extend(range(0x2000, 0x200A + 1)) + spu.extend(range(0x7F, 0xA0 + 1)) + + self.spaces = set([chr(i) for i in spu]) + + self.space_pat = re.compile(r'\s+', re.UNICODE) + + self.replace_char = { + u'`': u"'", + u'’': u"'", + u'´': u"'", + u'‘': u"'", + u'º': u'°', + u'–': u'-', + u'—': u'-' + } + + def sbc2dbc(self, ch): + n = ord(ch) + if 0xFF00 < n < 0xFF5F: + n -= 0xFEE0 + elif n == 0x3000: + n = 0x20 + else: + return ch + return chr(n) + + def clean(self, s): + try: + line = list(s.strip()) + size = len(line) + i = 0 + while i < size: + if line[i] < u' ' or line[i] in self.spaces: + line[i] = u' ' + else: + line[i] = self.replace_char.get(line[i], line[i]) + line[i] = self.sbc2dbc(line[i]) + + i += 1 + line = ''.join(line) + + line = self.space_pat.sub(' ', line).strip() + return line + except Exception: + return '' + + +if __name__ == '__main__': + + tc = TextClean() + + for line in sys.stdin: + res = tc.clean(line) + print(res) diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index 22421098..bd0f781c 100644 --- a/modelscope/utils/constant.py +++ b/modelscope/utils/constant.py @@ -176,6 +176,7 @@ class NLPTasks(object): relation_extraction = 'relation-extraction' zero_shot = 'zero-shot' translation = 'translation' + competency_aware_translation = 'competency-aware-translation' token_classification = 'token-classification' transformer_crf = 'transformer-crf' conversational = 'conversational' diff --git a/tests/pipelines/test_canmt_translation.py b/tests/pipelines/test_canmt_translation.py new file mode 100644 index 00000000..e3bce5d9 --- /dev/null +++ b/tests/pipelines/test_canmt_translation.py @@ -0,0 +1,68 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import unittest + +from modelscope.hub.snapshot_download import snapshot_download +from modelscope.models import Model +from modelscope.models.nlp import CanmtForTranslation +from modelscope.pipelines import pipeline +from modelscope.pipelines.nlp import CanmtTranslationPipeline +from modelscope.preprocessors import CanmtTranslationPreprocessor, Preprocessor +from modelscope.utils.constant import Tasks +from modelscope.utils.demo_utils import DemoCompatibilityCheck +from modelscope.utils.test_utils import test_level + + +class CanmtTranslationTest(unittest.TestCase, DemoCompatibilityCheck): + + def setUp(self) -> None: + self.task = Tasks.competency_aware_translation + self.model_id = 'damo/nlp_canmt_translation_zh2en_large' + + input = '110例癫痫患者血清抗脑抗体的测定' + input_2 = '世界是丰富多彩的。' + input_3 = '行业PE：处于PE估值历史分位较低的行业是房地产、纺织服饰、传媒。' + + @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') + def test_run_with_direct_download(self): + cache_path = snapshot_download(self.model_id) + preprocessor = Preprocessor.from_pretrained(cache_path) + pipeline1 = CanmtTranslationPipeline(cache_path, preprocessor) + pipeline2 = pipeline( + self.task, model=cache_path, preprocessor=preprocessor) + print( + f'pipeline1: {pipeline1(self.input)}\npipeline2: {pipeline2(self.input)}' + ) + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_run_with_model_name_batch(self): + run_kwargs = {'batch_size': 2} + pipeline_ins = pipeline(task=self.task, model=self.model_id) + print( + 'batch: ', + pipeline_ins([self.input, self.input_2, self.input_3], run_kwargs)) + + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') + def test_run_with_model_from_modelhub(self): + model = Model.from_pretrained(self.model_id) + preprocessor = Preprocessor.from_pretrained(model.model_dir) + pipeline_ins = pipeline( + task=self.task, model=model, preprocessor=preprocessor) + print(pipeline_ins(self.input)) + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_run_with_model_name(self): + pipeline_ins = pipeline(task=self.task, model=self.model_id) + print(pipeline_ins(self.input)) + + @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') + def test_run_with_default_model(self): + pipeline_ins = pipeline(task=self.task) + print(pipeline_ins(self.input)) + + @unittest.skip('demo compatibility test is only enabled on a needed-basis') + def test_demo_compatibility(self): + self.compatibility_check() + + +if __name__ == '__main__': + unittest.main() From bb879063e952707dce950910cf5189dba6fc9038 Mon Sep 17 00:00:00 2001 From: "pengteng.spt" Date: Fri, 7 Apr 2023 14:38:14 +0800 Subject: [PATCH 3/3] speech kws nearfield training add gradient accumulation config Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/12204960 --- .../trainers/audio/kws_nearfield_trainer.py | 1 + .../trainers/audio/kws_utils/batch_utils.py | 17 ++++++++++++----- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/modelscope/trainers/audio/kws_nearfield_trainer.py b/modelscope/trainers/audio/kws_nearfield_trainer.py index 9a84ce35..d1e3fdee 100644 --- a/modelscope/trainers/audio/kws_nearfield_trainer.py +++ b/modelscope/trainers/audio/kws_nearfield_trainer.py @@ -247,6 +247,7 @@ class KWSNearfieldTrainer(BaseTrainer): logger.info('Start training...') training_config = {} training_config['grad_clip'] = optim_conf['grad_clip'] + training_config['grad_accum'] = optim_conf.get('grad_accum', 1) training_config['log_interval'] = log_interval training_config['world_size'] = self.world_size training_config['rank'] = self.rank diff --git a/modelscope/trainers/audio/kws_utils/batch_utils.py b/modelscope/trainers/audio/kws_utils/batch_utils.py index ac382b79..75cf804e 100644 --- a/modelscope/trainers/audio/kws_utils/batch_utils.py +++ b/modelscope/trainers/audio/kws_utils/batch_utils.py @@ -44,6 +44,7 @@ def executor_train(model, optimizer, data_loader, device, writer, args): rank = args.get('rank', 0) local_rank = args.get('local_rank', 0) world_size = args.get('world_size', 1) + accum_batchs = args.get('grad_accum', 1) # [For distributed] Because iteration counts are not always equals between # processes, send stop-flag to the other processes if iterator is finished @@ -67,11 +68,16 @@ def executor_train(model, optimizer, data_loader, device, writer, args): logits, _ = model(feats) loss, acc = ctc_loss(logits, target, feats_lengths, target_lengths) loss = loss / num_utts - optimizer.zero_grad() + + # normlize loss to account for batch accumulation + loss = loss / accum_batchs loss.backward() - grad_norm = clip_grad_norm_(model.parameters(), clip) - if torch.isfinite(grad_norm): - optimizer.step() + if (batch_idx + 1) % accum_batchs == 0: + grad_norm = clip_grad_norm_(model.parameters(), clip) + if torch.isfinite(grad_norm): + optimizer.step() + optimizer.zero_grad() + if batch_idx % log_interval == 0: logger.info( 'RANK {}/{}/{} TRAIN Batch {}/{} size {} loss {:.6f}'.format( @@ -127,7 +133,8 @@ def executor_cv(model, data_loader, device, args): num_seen_tokens += target_lengths.sum() total_loss += loss.item() counter[0] += loss.item() - counter[1] += acc * target_lengths.sum() + counter[1] += acc * num_utts + # counter[1] += acc * target_lengths.sum() counter[2] += num_utts counter[3] += target_lengths.sum()