diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index 63ddd6d9..dd36bd47 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -405,6 +405,7 @@ class Pipelines(object): dialog_state_tracking = 'dialog-state-tracking' zero_shot_classification = 'zero-shot-classification' text_error_correction = 'text-error-correction' + word_alignment = 'word-alignment' plug_generation = 'plug-generation' gpt3_generation = 'gpt3-generation' gpt_moe_generation = 'gpt-moe-generation' @@ -925,6 +926,7 @@ class Preprocessors(object): sbert_token_cls_tokenizer = 'sbert-token-cls-tokenizer' zero_shot_cls_tokenizer = 'zero-shot-cls-tokenizer' text_error_correction = 'text-error-correction' + word_alignment = 'word-alignment' sentence_embedding = 'sentence-embedding' text_ranking = 'text-ranking' sequence_labeling_tokenizer = 'sequence-labeling-tokenizer' diff --git a/modelscope/models/nlp/bert/__init__.py b/modelscope/models/nlp/bert/__init__.py index 6578a0d7..00373f1c 100644 --- a/modelscope/models/nlp/bert/__init__.py +++ b/modelscope/models/nlp/bert/__init__.py @@ -17,6 +17,7 @@ if TYPE_CHECKING: from .token_classification import BertForTokenClassification from .document_segmentation import BertForDocumentSegmentation from .siamese_uie import SiameseUieModel + from .word_alignment import MBertForWordAlignment else: _import_structure = { 'backbone': [ diff --git a/modelscope/models/nlp/bert/word_alignment.py b/modelscope/models/nlp/bert/word_alignment.py new file mode 100644 index 00000000..526ec07d --- /dev/null +++ b/modelscope/models/nlp/bert/word_alignment.py @@ -0,0 +1,173 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team. +# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved. +# All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import torch.nn as nn +import torch.utils.checkpoint + +from modelscope.metainfo import Models +from modelscope.models.builder import MODELS +from modelscope.outputs import WordAlignmentOutput +from modelscope.utils import logger as logging +from modelscope.utils.constant import Tasks +from .backbone import BertModel, BertPreTrainedModel + +logger = logging.get_logger() + + +@MODELS.register_module(Tasks.word_alignment, module_name=Models.bert) +class MBertForWordAlignment(BertPreTrainedModel): + r"""MBert Model for the Word Alignment task. + + Code for EMNLP Findings 2022 paper, "Third-Party Aligner for Neural Word Alignments". + https://arxiv.org/abs/2211.04198 + + Parameters: + config (:class:`~modelscope.models.nlp.structbert.SbertConfig`): Model configuration class with + all the parameters of the model. + Initializing with a config file does not load the weights associated with the model, only the + configuration. Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model + weights. + """ + _keys_to_ignore_on_load_unexpected = [r'pooler'] + _keys_to_ignore_on_load_missing = [ + r'position_ids', r'predictions.decoder.bias' + ] + + def __init__(self, config, **kwargs): + super().__init__(config) + + if config.is_decoder: + logger.warning( + 'If you want to use `BertForMaskedLM` make sure `config.is_decoder=False` for ' + 'bi-directional self-attention.') + config.num_hidden_layers = kwargs.get('encoder_layers', 8) + + self.bert = BertModel(config, add_pooling_layer=False) + + # Initialize weights and apply final processing + self.post_init() + + def forward( + self, + src_input_ids=None, + src_attention_mask=None, + src_b2w_map=None, + tgt_input_ids=None, + tgt_attention_mask=None, + tgt_b2w_map=None, + threshold=0.001, + bpe_level=False, + ): + """ + Args: src_input_ids: + Indices of source input sequence tokens in the vocabulary. + src_attention_mask: + Source mask to avoid performing attention on padding token indices. + src_b2w_map: + Word order numner of subword in source sequence. + tgt_input_ids: + Indices of target input sequence tokens in the vocabulary. + tgt_attention_mask: + Target mask to avoid performing attention on padding token indices. + tgt_b2w_map: + Word order numner of subword in target sequence. + threshold: + The threshold used to extract alignment. + bpe_level: + Return subword-level alignment or not. + Example: + { + 'src_input_ids': LongTensor([[2478,242,24,4]]), + 'src_attention_mask': BoolTensor([[1,1,1,1]]), + 'src_b2w_map': LongTensor([[0,1,2,3]]), + 'tgt_input_ids': LongTensor([[1056,356,934,263,7]]), + 'tgt_attention_mask': BoolTensor([[1,1,1,1,1]]), + 'tgt_b2w_map': longtensor([[0,1,1,2,3]]), + 'threshold': 0.001, + 'bpe_level': False, + } + Returns `modelscope.outputs.WordAlignmentOutput` + """ + with torch.no_grad(): + src_encoder_out = self.bert( + input_ids=src_input_ids, + attention_mask=src_attention_mask.float(), + head_mask=None, + inputs_embeds=None, + output_hidden_states=True, + ) + tgt_encoder_out = self.bert( + input_ids=tgt_input_ids, + attention_mask=tgt_attention_mask.float(), + head_mask=None, + inputs_embeds=None, + output_hidden_states=True, + ) + + atten_mask_src = (1 - ( + (src_input_ids != 101) & (src_input_ids != 102) + & src_attention_mask)[:, None, None, :].float()) * -10000 + atten_mask_tgt = (1 - ( + (tgt_input_ids != 101) & (tgt_input_ids != 102) + & tgt_attention_mask)[:, None, None, :].float()) * -10000 + + src_align_out = src_encoder_out[0] + tgt_align_out = tgt_encoder_out[0] + + bpe_sim = torch.bmm(src_align_out, tgt_align_out.transpose(1, 2)) + + attention_scores_src = bpe_sim.unsqueeze(1) + atten_mask_tgt + attention_scores_tgt = bpe_sim.unsqueeze(1) + atten_mask_src.transpose( + -1, -2) + + attention_probs_src = nn.Softmax(dim=-1)(attention_scores_src) + attention_probs_tgt = nn.Softmax(dim=-2)(attention_scores_tgt) + + align_matrix = (attention_probs_src > threshold) * ( + attention_probs_tgt > threshold) + align_matrix = align_matrix.squeeze(1) + + len_src = (atten_mask_src == 0).sum(dim=-1).unsqueeze(-1) + len_tgt = (atten_mask_tgt == 0).sum(dim=-1).unsqueeze(-1) + + attention_probs_src = nn.Softmax(dim=-1)( + attention_scores_src / torch.sqrt(len_src.float())) + attention_probs_tgt = nn.Softmax(dim=-2)( + attention_scores_tgt / torch.sqrt(len_tgt.float())) + + word_aligns = [] + + for idx, (line_align, b2w_src, b2w_tgt) in enumerate( + zip(align_matrix, src_b2w_map, tgt_b2w_map)): + aligns = dict() + non_specials = torch.where(line_align) + for i, j in zip(*non_specials): + if not bpe_level: + word_pair = (src_b2w_map[idx][i - 1].item(), + tgt_b2w_map[idx][j - 1].item()) + if word_pair not in aligns: + aligns[word_pair] = bpe_sim[idx][i, j].item() + else: + aligns[word_pair] = max(aligns[word_pair], + bpe_sim[idx][i, j].item()) + else: + aligns[(i.item() - 1, + j.item() - 1)] = bpe_sim[idx][i, j].item() + word_aligns.append(aligns) + + return WordAlignmentOutput(predictions=word_aligns) diff --git a/modelscope/outputs/nlp_outputs.py b/modelscope/outputs/nlp_outputs.py index 8ed8c94a..a48e3b0e 100644 --- a/modelscope/outputs/nlp_outputs.py +++ b/modelscope/outputs/nlp_outputs.py @@ -1,5 +1,5 @@ from dataclasses import dataclass -from typing import Optional, Tuple, Union +from typing import List, Optional, Tuple, Union import numpy as np @@ -328,6 +328,14 @@ class TextErrorCorrectionOutput(ModelOutputBase): predictions: np.ndarray = None +@dataclass +class WordAlignmentOutput(ModelOutputBase): + """The output class for word alignment models. + """ + + predictions: List = None + + @dataclass class TextGenerationModelOutput(ModelOutputBase): """The output class for text generation models. diff --git a/modelscope/outputs/outputs.py b/modelscope/outputs/outputs.py index d8217aa0..75d9eaa5 100644 --- a/modelscope/outputs/outputs.py +++ b/modelscope/outputs/outputs.py @@ -709,6 +709,11 @@ TASK_OUTPUTS = { # "output": "我想吃苹果" # } Tasks.text_error_correction: [OutputKeys.OUTPUT], + # word_alignment result for a single sample + # { + # "output": "0-0 1-3 2-4 3-1 4-2 5-5" + # } + Tasks.word_alignment: [OutputKeys.OUTPUT], Tasks.sentence_embedding: [OutputKeys.TEXT_EMBEDDING, OutputKeys.SCORES], Tasks.text_ranking: [OutputKeys.SCORES], diff --git a/modelscope/pipelines/nlp/__init__.py b/modelscope/pipelines/nlp/__init__.py index b4ea93be..55c98776 100644 --- a/modelscope/pipelines/nlp/__init__.py +++ b/modelscope/pipelines/nlp/__init__.py @@ -25,6 +25,7 @@ if TYPE_CHECKING: from .summarization_pipeline import SummarizationPipeline from .translation_quality_estimation_pipeline import TranslationQualityEstimationPipeline from .text_error_correction_pipeline import TextErrorCorrectionPipeline + from .word_alignment_pipeline import WordAlignmentPipeline from .text_generation_pipeline import TextGenerationPipeline, TextGenerationT5Pipeline from .fid_dialogue_pipeline import FidDialoguePipeline from .token_classification_pipeline import TokenClassificationPipeline @@ -70,6 +71,7 @@ else: ['TableQuestionAnsweringPipeline'], 'text_classification_pipeline': ['TextClassificationPipeline'], 'text_error_correction_pipeline': ['TextErrorCorrectionPipeline'], + 'word_alignment_pipeline': ['WordAlignmentPipeline'], 'text_generation_pipeline': ['TextGenerationPipeline', 'TextGenerationT5Pipeline'], 'fid_dialogue_pipeline': ['FidDialoguePipeline'], diff --git a/modelscope/pipelines/nlp/word_alignment_pipeline.py b/modelscope/pipelines/nlp/word_alignment_pipeline.py new file mode 100644 index 00000000..8513f28e --- /dev/null +++ b/modelscope/pipelines/nlp/word_alignment_pipeline.py @@ -0,0 +1,68 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +from typing import Any, Dict, Optional, Union + +import numpy as np + +from modelscope.metainfo import Pipelines +from modelscope.models import Model +from modelscope.outputs import OutputKeys +from modelscope.pipelines.base import Pipeline +from modelscope.pipelines.builder import PIPELINES +from modelscope.preprocessors import WordAlignmentPreprocessor +from modelscope.utils.constant import Tasks + +__all__ = ['WordAlignmentPipeline'] + + +@PIPELINES.register_module( + Tasks.word_alignment, module_name=Pipelines.word_alignment) +class WordAlignmentPipeline(Pipeline): + + def __init__(self, + model: Union[Model, str], + preprocessor: WordAlignmentPreprocessor = None, + config_file: str = None, + device: str = 'gpu', + auto_collate=True, + sequence_length=128, + **kwargs): + """Use `model` and `preprocessor` to create a nlp text dual encoder then generates the text representation. + Args: + model (str or Model): Supply either a local model dir which supported the WS task, + or a model id from the model hub, or a torch model instance. + preprocessor (Preprocessor): A WordAlignmentPreprocessor. + kwargs (dict, `optional`): + Extra kwargs passed into the preprocessor's constructor. + Example: + >>> from modelscope.pipelines import pipeline + >>> from modelscope.utils.constant import Tasks + >>> model_id = 'damo/Third-Party-Supervised-Word-Aligner-mBERT-base-zhen' + >>> input = {"sentence_pair": '贝利 在 墨西哥 推出 自传 。||| pele promotes autobiography in mexico .'} + >>> pipeline_ins = pipeline(Tasks.word_alignment, model=model_id) + >>> print(pipeline_ins(input)['output']) + """ + super().__init__( + model=model, + preprocessor=preprocessor, + config_file=config_file, + device=device, + auto_collate=auto_collate) + if preprocessor is None: + self.preprocessor = WordAlignmentPreprocessor.from_pretrained( + self.model.model_dir, + sequence_length=sequence_length, + **kwargs) + + def forward(self, inputs: Dict[str, Any], + **forward_params) -> Dict[str, Any]: + return self.model(**inputs, **forward_params) + + def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: + + align = [] + for k in inputs[0][0].keys(): + align.append(f'{k[0]}-{k[1]}') + align = ' '.join(align) + + return {OutputKeys.OUTPUT: align} diff --git a/modelscope/preprocessors/__init__.py b/modelscope/preprocessors/__init__.py index 3a13828b..1cd25e50 100644 --- a/modelscope/preprocessors/__init__.py +++ b/modelscope/preprocessors/__init__.py @@ -30,9 +30,9 @@ if TYPE_CHECKING: TextGenerationSentencePiecePreprocessor, TokenClassificationTransformersPreprocessor, TextErrorCorrectionPreprocessor, TextGenerationT5Preprocessor, - TextGenerationTransformersPreprocessor, Tokenize, - WordSegmentationBlankSetToLabelPreprocessor, CodeGeeXPreprocessor, - MGLMSummarizationPreprocessor, + WordAlignmentPreprocessor, TextGenerationTransformersPreprocessor, + Tokenize, WordSegmentationBlankSetToLabelPreprocessor, + CodeGeeXPreprocessor, MGLMSummarizationPreprocessor, ZeroShotClassificationTransformersPreprocessor, TextGenerationJiebaPreprocessor, SentencePiecePreprocessor, DialogIntentPredictionPreprocessor, DialogModelingPreprocessor, diff --git a/modelscope/preprocessors/nlp/__init__.py b/modelscope/preprocessors/nlp/__init__.py index 23fa9f94..f0660374 100644 --- a/modelscope/preprocessors/nlp/__init__.py +++ b/modelscope/preprocessors/nlp/__init__.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING from modelscope.utils.import_utils import LazyImportModule if TYPE_CHECKING: + from .word_alignment_preprocessor import WordAlignmentPreprocessor from .text_error_correction import TextErrorCorrectionPreprocessor from .text_generation_preprocessor import TextGenerationJiebaPreprocessor from .bert_seq_cls_tokenizer import Tokenize @@ -66,6 +67,9 @@ else: 'text_error_correction': [ 'TextErrorCorrectionPreprocessor', ], + 'word_alignment_preprocessor': [ + 'WordAlignmentPreprocessor', + ], 'mglm_summarization_preprocessor': ['MGLMSummarizationPreprocessor'], 'token_classification_thai_preprocessor': [ 'NERPreprocessorThai', diff --git a/modelscope/preprocessors/nlp/word_alignment_preprocessor.py b/modelscope/preprocessors/nlp/word_alignment_preprocessor.py new file mode 100644 index 00000000..390deb40 --- /dev/null +++ b/modelscope/preprocessors/nlp/word_alignment_preprocessor.py @@ -0,0 +1,131 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import itertools +import os +import os.path as osp +from typing import Any, Dict, Optional, Union + +import numpy as np +import torch + +from modelscope.metainfo import Preprocessors +from modelscope.preprocessors.base import Preprocessor +from modelscope.preprocessors.builder import PREPROCESSORS +from modelscope.utils.constant import Fields, ModeKeys +from modelscope.utils.hub import get_model_type +from modelscope.utils.logger import get_logger +from .transformers_tokenizer import NLPTokenizer + + +@PREPROCESSORS.register_module( + Fields.nlp, module_name=Preprocessors.word_alignment) +class WordAlignmentPreprocessor(Preprocessor): + """The tokenizer preprocessor used in word alignment . + """ + + def __init__(self, + model_dir: str, + sequence_pair='sentence_pair', + mode=ModeKeys.INFERENCE, + use_fast: bool = False, + sequence_length: int = None, + **kwargs): + """The preprocessor for word alignment task. + + Args: + model_dir: The model dir used to initialize the tokenizer. + sequence_pair: The key of the sequence pair. + mode: The mode for the preprocessor. + use_fast: Use the fast tokenizer or not. + sequence_length: The max sequence length which the model supported, + will be passed into tokenizer as the 'max_length' param. + **kwargs: Extra args input. + {sequence_length: The sequence length which the model supported.} + """ + self.sequence_pair = sequence_pair + + kwargs[ + 'sequence_length'] = sequence_length if sequence_length is not None else kwargs.get( + 'max_length', 128) + self.max_length = kwargs['sequence_length'] + kwargs.pop('max_length', None) + model_type = None + + if model_dir is not None: + model_type = get_model_type(model_dir) + self.nlp_tokenizer = NLPTokenizer( + model_dir, model_type, use_fast=use_fast, tokenize_kwargs=kwargs) + super().__init__(mode=mode) + + def __call__(self, data: Dict, **kwargs) -> Dict[str, Any]: + """process the raw input data + + Args: + data Dict: + Example: + {"sentence_pair": "贝利 在 墨西哥 推出 自传 。||| pele promotes autobiography in mexico ."} + Returns: + Dict[str, Any]: the preprocessed data + """ + sentence_pair = data[self.sequence_pair] + source_sentences, target_sentences = sentence_pair.split('|||') + # src_lang = data.get("src_lang", 'en_XX') + # tgt_lang = data.get("tgt_lang", 'en_XX') + if 'return_tensors' not in kwargs: + kwargs[ + 'return_tensors'] = 'pt' if self.mode == ModeKeys.INFERENCE else None + + sent_src, sent_tgt = source_sentences.strip().split( + ), target_sentences.strip().split() + + token_src = [ + self.nlp_tokenizer.tokenizer.tokenize(word) for word in sent_src + ] + token_tgt = [ + self.nlp_tokenizer.tokenizer.tokenize(word) for word in sent_tgt + ] + wid_src = [ + self.nlp_tokenizer.tokenizer.convert_tokens_to_ids(x) + for x in token_src + ] + wid_tgt = [ + self.nlp_tokenizer.tokenizer.convert_tokens_to_ids(x) + for x in token_tgt + ] + + ids_tgt = self.nlp_tokenizer.tokenizer.prepare_for_model( + list(itertools.chain(*wid_tgt)), + return_tensors='pt', + max_length=self.max_length, + prepend_batch_axis=True)['input_ids'] + ids_src = self.nlp_tokenizer.tokenizer.prepare_for_model( + list(itertools.chain(*wid_src)), + return_tensors='pt', + max_length=self.max_length, + prepend_batch_axis=True)['input_ids'] + + bpe2word_map_src = [] + for i, word_list in enumerate(token_src): + bpe2word_map_src += [i for x in word_list] + bpe2word_map_src = torch.Tensor(bpe2word_map_src).type_as( + ids_src).view(1, -1) + bpe2word_map_tgt = [] + for i, word_list in enumerate(token_tgt): + bpe2word_map_tgt += [i for x in word_list] + bpe2word_map_tgt = torch.Tensor(bpe2word_map_tgt).type_as( + ids_tgt).view(1, -1) + attention_mask_src = ( + ids_src != self.nlp_tokenizer.tokenizer.pad_token_id) + attention_mask_tgt = ( + ids_tgt != self.nlp_tokenizer.tokenizer.pad_token_id) + + return { + 'src_input_ids': ids_src, + 'src_attention_mask': attention_mask_src, + 'src_b2w_map': bpe2word_map_src, + 'tgt_input_ids': ids_tgt, + 'tgt_attention_mask': attention_mask_tgt, + 'tgt_b2w_map': bpe2word_map_tgt, + 'threshold': 0.001, + 'bpe_level': False + } diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index c2d0c6f5..f289ea50 100644 --- a/modelscope/utils/constant.py +++ b/modelscope/utils/constant.py @@ -186,6 +186,7 @@ class NLPTasks(object): zero_shot_classification = 'zero-shot-classification' backbone = 'backbone' text_error_correction = 'text-error-correction' + word_alignment = 'word-alignment' faq_question_answering = 'faq-question-answering' information_extraction = 'information-extraction' document_segmentation = 'document-segmentation'