Add word alignment model (#101)

* word alignment

* Delete ast_index_file.py

* add introduction

* fix-wordalignmentpreprocessor

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'

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': [

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)

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.

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],
 

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'],

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}

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,

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',

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
+        }

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'