From 81a75bf260694c71380ff2cc7c92bd643d6aa295 Mon Sep 17 00:00:00 2001 From: "hehong.chh" Date: Thu, 9 Feb 2023 08:11:06 +0000 Subject: [PATCH] [to #42322933] add fuse-in-decoder dialogue task --- modelscope/metainfo.py | 2 + modelscope/models/nlp/fid_plug/__init__.py | 36 + modelscope/models/nlp/fid_plug/backbone.py | 1148 +++++++++++++++++ .../models/nlp/fid_plug/configuration.py | 103 ++ .../models/nlp/fid_plug/text_generation.py | 180 +++ modelscope/outputs/outputs.py | 6 + modelscope/pipeline_inputs.py | 6 + modelscope/pipelines/nlp/__init__.py | 2 + .../pipelines/nlp/fid_dialogue_pipeline.py | 176 +++ .../nlp/transformers_tokenizer.py | 2 +- modelscope/utils/constant.py | 1 + tests/pipelines/test_plug_dialogue.py | 45 + 12 files changed, 1706 insertions(+), 1 deletion(-) create mode 100644 modelscope/models/nlp/fid_plug/__init__.py create mode 100644 modelscope/models/nlp/fid_plug/backbone.py create mode 100644 modelscope/models/nlp/fid_plug/configuration.py create mode 100644 modelscope/models/nlp/fid_plug/text_generation.py create mode 100644 modelscope/pipelines/nlp/fid_dialogue_pipeline.py create mode 100644 tests/pipelines/test_plug_dialogue.py diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index d156819b..7350ac4d 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -130,6 +130,7 @@ class Models(object): unite = 'unite' megatron_bert = 'megatron-bert' use = 'user-satisfaction-estimation' + fid_plug = 'fid-plug' plug_mental = 'plug-mental' # audio models @@ -339,6 +340,7 @@ class Pipelines(object): named_entity_recognition_thai = 'named-entity-recognition-thai' named_entity_recognition_viet = 'named-entity-recognition-viet' text_generation = 'text-generation' + fid_dialogue = 'fid-dialogue' text2text_generation = 'text2text-generation' sentiment_analysis = 'sentiment-analysis' sentiment_classification = 'sentiment-classification' diff --git a/modelscope/models/nlp/fid_plug/__init__.py b/modelscope/models/nlp/fid_plug/__init__.py new file mode 100644 index 00000000..fe455e84 --- /dev/null +++ b/modelscope/models/nlp/fid_plug/__init__.py @@ -0,0 +1,36 @@ +# Copyright 2021-2022 The Alibaba DAMO NLP Team Authors. +# 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. +from typing import TYPE_CHECKING + +from modelscope.utils.import_utils import LazyImportModule + +if TYPE_CHECKING: + from .configuration import PlugConfig + from .text_generation import (PlugV2Chat, PlugV2FidChat) +else: + _import_structure = { + 'configuration': ['PlugConfig'], + 'text_generation': ['PlugV2Chat', 'PlugV2FidChat'], + } + + import sys + + sys.modules[__name__] = LazyImportModule( + __name__, + globals()['__file__'], + _import_structure, + module_spec=__spec__, + extra_objects={}, + ) diff --git a/modelscope/models/nlp/fid_plug/backbone.py b/modelscope/models/nlp/fid_plug/backbone.py new file mode 100644 index 00000000..e3e91606 --- /dev/null +++ b/modelscope/models/nlp/fid_plug/backbone.py @@ -0,0 +1,1148 @@ +# Copyright 2021-2022 The Alibaba DAMO NLP Team Authors. +# Copyright 2018 The Google AI Language Team Authors and The HugginFace Inc. team. +# Copyright (c) 2019, NVIDIA CORPORATION. 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 copy +import math +import os +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Union + +import numpy as np +import torch +import torch.nn.functional as F +from torch import Tensor, nn +from torch.nn.init import xavier_uniform_ +from transformers import (BertConfig, BertModel, BertTokenizer, RobertaConfig, + RobertaModel, RobertaTokenizer, logging) +from transformers.activations import ACT2FN +from transformers.modeling_utils import PreTrainedModel + +from .configuration import PlugConfig + +CONFIG_NAME = 'config.json' +WEIGHTS_NAME = 'pytorch_model.bin' + + +class MultiHeadedAttention(nn.Module): # SelfAttention + """ + Multi-Head Attention module from + "Attention is All You Need" + :cite:`DBLP:journals/corr/VaswaniSPUJGKP17`. + + Similar to standard `dot` attention but uses + multiple attention distributions simulataneously + to select relevant items. + + .. mermaid:: + + graph BT + A[key] + B[value] + C[query] + O[output] + subgraph Attn + D[Attn 1] + E[Attn 2] + F[Attn N] + end + A --> D + C --> D + A --> E + C --> E + A --> F + C --> F + D --> O + E --> O + F --> O + B --> O + + Also includes several additional tricks. + + Args: + head_count (int): number of parallel heads + model_dim (int): the dimension of keys/values/queries, + must be divisible by head_count + dropout (float): dropout parameter + """ + + def __init__(self, + head_count, + model_dim, + dropout=0.1, + use_final_linear=True): + assert model_dim % head_count == 0 + self.dim_per_head = model_dim // head_count + self.model_dim = model_dim + + super().__init__() + self.head_count = head_count + + self.linear_keys = nn.Linear(model_dim, head_count * self.dim_per_head) + self.linear_values = nn.Linear(model_dim, + head_count * self.dim_per_head) + self.linear_query = nn.Linear(model_dim, + head_count * self.dim_per_head) + self.softmax = nn.Softmax(dim=-1) + self.dropout = nn.Dropout(dropout) + self.use_final_linear = use_final_linear + if (self.use_final_linear): + self.final_linear = nn.Linear(model_dim, model_dim) + + def forward(self, + key, + value, + query, + mask=None, + layer_cache=None, + type=None, + predefined_graph_1=None, + return_attn=False): + """ + Compute the context vector and the attention vectors. + + Args: + key (`FloatTensor`): set of `key_len` + key vectors `[batch, key_len, dim]` + value (`FloatTensor`): set of `key_len` + value vectors `[batch, key_len, dim]` + query (`FloatTensor`): set of `query_len` + query vectors `[batch, query_len, dim]` + mask: binary mask indicating which keys have + non-zero attention `[batch, query_len, key_len]` + Returns: + (`FloatTensor`, `FloatTensor`) : + + * output context vectors `[batch, query_len, dim]` + * one of the attention vectors `[batch, query_len, key_len]` + """ + + batch_size = key.size(0) + dim_per_head = self.dim_per_head + head_count = self.head_count + + def shape(x): + """ projection """ + return x.view(batch_size, -1, head_count, dim_per_head) \ + .transpose(1, 2) + + def unshape(x): + """ compute context """ + return x.transpose(1, 2).contiguous() \ + .view(batch_size, -1, head_count * dim_per_head) + + # 1) Project key, value, and query. + if layer_cache is not None: + if type == 'self': + query, key, value = self.linear_query(query), self.linear_keys( + query), self.linear_values(query) + + key = shape(key) + value = shape(value) + + device = key.device + if layer_cache['self_keys'] is not None: + key = torch.cat((layer_cache['self_keys'].to(device), key), + dim=2) + if layer_cache['self_values'] is not None: + value = torch.cat( + (layer_cache['self_values'].to(device), value), dim=2) + layer_cache['self_keys'] = key + layer_cache['self_values'] = value + elif type == 'context': + query = self.linear_query(query) + if layer_cache['memory_keys'] is None: + key, value = self.linear_keys(key), self.linear_values( + value) + key = shape(key) + value = shape(value) + else: + key, value = layer_cache['memory_keys'], layer_cache[ + 'memory_values'] + layer_cache['memory_keys'] = key + layer_cache['memory_values'] = value + else: + key = self.linear_keys(key) + value = self.linear_values(value) + query = self.linear_query(query) + key = shape(key) + value = shape(value) + + query = shape(query) + + # 2) Calculate and scale scores. + query = query / math.sqrt(dim_per_head) + scores = torch.matmul(query, key.transpose(2, 3)) + + if mask is not None: + mask = mask.unsqueeze(1).expand_as(scores) + scores = scores.masked_fill(mask, float('-inf')) + + # 3) Apply attention dropout and compute context vectors. + + attn = self.softmax(scores) + + if (predefined_graph_1 is not None): + attn_masked = attn[:, -1] * predefined_graph_1 + attn_masked = attn_masked / ( + torch.sum(attn_masked, 2).unsqueeze(2) + 1e-9) + + attn = torch.cat([attn[:, :-1], attn_masked.unsqueeze(1)], 1) + + drop_attn = self.dropout(attn) + if (self.use_final_linear): + context = unshape(torch.matmul(drop_attn, value)) + output = self.final_linear(context) + if return_attn: + return output, attn + else: + return output + else: + context = torch.matmul(drop_attn, value) + if return_attn: + return context, attn + else: + return context + + +class PositionwiseFeedForward(nn.Module): # Output + """ A two-layer Feed-Forward-Network with residual layer norm. + + Args: + d_model (int): the size of input for the first-layer of the FFN. + d_ff (int): the hidden layer size of the second-layer + of the FNN. + dropout (float): dropout probability in :math:`[0, 1)`. + """ + + def __init__(self, d_model, d_ff, dropout=0.1): + super().__init__() + self.layer_norm = nn.LayerNorm(d_model, eps=1e-6) + self.w_1 = nn.Linear(d_model, d_ff) + self.actv = ACT2FN['gelu_new'] + self.dropout_1 = nn.Dropout(dropout) + self.w_2 = nn.Linear(d_ff, d_model) + self.dropout_2 = nn.Dropout(dropout) + + def forward(self, x): + inter = self.dropout_1(self.actv(self.w_1(self.layer_norm(x)))) + output = self.dropout_2(self.w_2(inter)) + return output + x + + +class TransformerDecoderLayer(nn.Module): # Layer + """ + Args: + d_model (int): the dimension of keys/values/queries in + MultiHeadedAttention, also the input size of + the first-layer of the PositionwiseFeedForward. + heads (int): the number of heads for MultiHeadedAttention. + d_ff (int): the second-layer of the PositionwiseFeedForward. + dropout (float): dropout probability(0-1.0). + self_attn_type (string): type of self-attention scaled-dot, average + """ + MAX_SIZE = 5000 + + def __init__(self, d_model, heads, d_ff, dropout): + super().__init__() + + self.self_attn = MultiHeadedAttention(heads, d_model, dropout=dropout) + + self.context_attn = MultiHeadedAttention( + heads, d_model, dropout=dropout) + self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout) + self.layer_norm_1 = nn.LayerNorm(d_model, eps=1e-6) + self.layer_norm_2 = nn.LayerNorm(d_model, eps=1e-6) + self.drop = nn.Dropout(dropout) + mask = self._get_attn_subsequent_mask(self.MAX_SIZE) + # Register self.mask as a buffer in TransformerDecoderLayer, so + # it gets TransformerDecoderLayer's cuda behavior automatically. + self.register_buffer('mask', mask) + + def forward(self, + inputs, + memory_bank, + src_pad_mask, + tgt_pad_mask, + previous_input=None, + layer_cache=None, + step=None): + """ + Args: + inputs (`FloatTensor`): `[batch_size x 1 x model_dim]` + memory_bank (`FloatTensor`): `[batch_size x src_len x model_dim]` + src_pad_mask (`LongTensor`): `[batch_size x 1 x src_len]` + tgt_pad_mask (`LongTensor`): `[batch_size x 1 x 1]` + + Returns: + (`FloatTensor`, `FloatTensor`, `FloatTensor`): + + * output `[batch_size x 1 x model_dim]` + * attn `[batch_size x 1 x src_len]` + * all_input `[batch_size x current_step x model_dim]` + + """ + dec_mask = torch.gt( + tgt_pad_mask.type(torch.uint8) + + self.mask[:, :tgt_pad_mask.size(1), :tgt_pad_mask.size(1)].type( + torch.uint8), 0) + input_norm = self.layer_norm_1(inputs) + all_input = input_norm + if previous_input is not None: + all_input = torch.cat((previous_input, input_norm), dim=1) + dec_mask = None + + query = self.self_attn( + all_input, + all_input, + input_norm, + mask=dec_mask, + layer_cache=layer_cache, + type='self') + + query = self.drop(query) + inputs + + query_norm = self.layer_norm_2(query) + mid, attn = self.context_attn( + memory_bank, + memory_bank, + query_norm, + mask=src_pad_mask, + layer_cache=layer_cache, + type='context', + return_attn=True) + output = self.feed_forward(self.drop(mid) + query) + + return output, attn, all_input + + def _get_attn_subsequent_mask(self, size): + """ + Get an attention mask to avoid using the subsequent info. + + Args: + size: int + + Returns: + (`LongTensor`): + + * subsequent_mask `[1 x size x size]` + """ + attn_shape = (1, size, size) + subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8') + subsequent_mask = torch.from_numpy(subsequent_mask) + return subsequent_mask + + +class PositionalEncoding(nn.Module): + + def __init__(self, dropout, dim, max_len=5000): + super().__init__() + pe = torch.zeros(max_len, dim) + position = torch.arange(0, max_len).unsqueeze(1) + div_term = torch.exp((torch.arange(0, dim, 2, dtype=torch.float) + * -(math.log(10000.0) / dim))) + pe[:, 0::2] = torch.sin(position.float() * div_term) + pe[:, 1::2] = torch.cos(position.float() * div_term) + pe = pe.unsqueeze(0) + self.register_buffer('pe', pe) + self.dropout = nn.Dropout(dropout) + self.dim = dim + + def forward(self, emb, step=None): + emb = emb * math.sqrt(self.dim) + if (step): + emb = emb + self.pe[:, step][:, None, :] + + else: + emb = emb + self.pe[:, :emb.size(1)] + emb = self.dropout(emb) + return emb + + def get_emb(self, emb): + return self.pe[:, :emb.size(1)] + + +class TransformerDecoderState: + + def __init__(self, src: Tensor, cache_num_layers: int = -1): + self.src: Tensor = src + self.previous_input: Tensor = None + self.previous_layer_inputs: Tensor = None + self.cache: Optional[Dict[str, Any]] = None + if cache_num_layers != -1: + self._init_cache(cache_num_layers) + + def update_state(self, new_input, previous_layer_inputs): + self.previous_input = new_input + self.previous_layer_inputs = previous_layer_inputs + self.cache = None + + def _init_cache(self, num_layers): + self.cache = {} + for layer in range(num_layers): + layer_cache = {'memory_keys': None, 'memory_values': None} + layer_cache['self_keys'] = None + layer_cache['self_values'] = None + self.cache['layer_{}'.format(layer)] = layer_cache + + def map_batch_fn(self, fn): + + def _recursive_map(struct, batch_dim=0): + for k, v in struct.items(): + if v is not None: + if isinstance(v, dict): + _recursive_map(v) + else: + struct[k] = fn(v, batch_dim) + + self.src = fn(self.src, 0) + if self.cache is not None: + _recursive_map(self.cache) + + +class TransformerDecoder(nn.Module): # Decoder + """ + The Transformer decoder from "Attention is All You Need". + + + .. mermaid:: + + graph BT + A[input] + B[multi-head self-attn] + BB[multi-head src-attn] + C[feed forward] + O[output] + A --> B + B --> BB + BB --> C + C --> O + + + Args: + num_layers (int): number of encoder layers. + d_model (int): size of the model + heads (int): number of heads + d_ff (int): size of the inner FF layer + dropout (float): dropout parameters + embeddings (:obj:`onmt.modules.Embeddings`): + embeddings to use, should have positional encodings + attn_type (str): if using a seperate copy attention + """ + decoder_type = 'transformer' + + def __init__(self, num_layers, d_model, heads, d_ff, dropout, embeddings): + super().__init__() + + # Basic attributes. + self.num_layers = num_layers + self.embeddings = embeddings + self.pos_emb = PositionalEncoding(dropout, + self.embeddings.embedding_dim) + + # Build TransformerDecoder. + self.transformer_layers = nn.ModuleList([ + TransformerDecoderLayer(d_model, heads, d_ff, dropout) + for _ in range(num_layers) + ]) + self.layer_norm = nn.LayerNorm(d_model, eps=1e-6) + self.state = None + + def forward(self, + state: TransformerDecoderState, + tgt: Tensor, + memory_bank: Tensor, + step: int = None, + memory_masks: Tensor = None): + src_words = state.src + tgt_words = tgt + src_batch, src_len = src_words.size() + tgt_batch, tgt_len = tgt_words.size() + + # Run the forward pass of the TransformerDecoder. + # emb = self.embeddings(tgt, step=step) + emb = self.embeddings(tgt) + assert emb.dim() == 3 # len x batch x embedding_dim + output = self.pos_emb(emb, step) + + src_memory_bank = memory_bank + padding_idx = self.embeddings.padding_idx + tgt_pad_mask = tgt_words.data.eq(padding_idx).unsqueeze(1) \ + .expand(tgt_batch, tgt_len, tgt_len) + + if (memory_masks is not None): + src_len = memory_masks.size(-1) + src_pad_mask = memory_masks.expand(src_batch, tgt_len, src_len) + else: + src_pad_mask = src_words.data.eq(padding_idx).unsqueeze(1) \ + .expand(src_batch, tgt_len, src_len) + + if state.cache is None: + saved_inputs = [] + attns = [] + for i in range(self.num_layers): + prev_layer_input = None + if state.cache is None: + if state.previous_input is not None: + prev_layer_input = state.previous_layer_inputs[i] + output, attn, all_input \ + = self.transformer_layers[i]( + output, src_memory_bank, + src_pad_mask, tgt_pad_mask, + previous_input=prev_layer_input, + layer_cache=state.cache['layer_{}'.format(i)] + if state.cache is not None else None, + step=step) + if state.cache is None: + saved_inputs.append(all_input) + attns.append(attn) + + if state.cache is None: + saved_inputs = torch.stack(saved_inputs) + + output = self.layer_norm(output) + + # Process the result and update the attentions. + if state.cache is None: + state.update_state(tgt, saved_inputs) + + return output, attns, state + + +class PlugPointerGenerator(nn.Module): + + def __init__(self, hidden_size, vocab_size): + super().__init__() + self.dense = nn.Linear(hidden_size, vocab_size) + self.gen_func = nn.LogSoftmax(-1) + + def forward(self, x): + x = self.dense(x) + x = self.gen_func(x) + return x + + +class PlugPreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = PlugConfig + base_model_prefix = 'plug' + + @classmethod + def from_pretrained( + cls, pretrained_model_name_or_path: Optional[Union[str, + os.PathLike]]): + config_file = os.path.join(pretrained_model_name_or_path, CONFIG_NAME) + config = PlugConfig.from_json_file(config_file) if os.path.isfile( + config_file) else PlugConfig() + config.encoder_pth = os.path.join(pretrained_model_name_or_path, + config.encoder_pth) + checkpoint_file = os.path.join(pretrained_model_name_or_path, + WEIGHTS_NAME) + checkpoint = torch.load(checkpoint_file) if os.path.isfile( + checkpoint_file) else None + return cls(config, checkpoint) + + +class PlugModel(PlugPreTrainedModel): # Model + + def __init__(self, config, checkpoint=None): + super().__init__(config) + self.config = config + if config.encoder == 'bert' or config.encoder == 'zh_bert': + self.bert = BertModel( + BertConfig.from_pretrained(config.encoder_pth)) + elif config.encoder == 'roberta': + self.bert = RobertaModel( + RobertaConfig.from_pretrained(config.encoder_pth)) + + if (config.max_pos > 512): + my_pos_embeddings = nn.Embedding( + config.max_pos, self.bert.model.config.hidden_size) + my_pos_embeddings.weight.data[: + 512] = self.bert.embeddings.position_embeddings.weight.data + my_pos_embeddings.weight.data[ + 512:] = self.bert.embeddings.position_embeddings.weight.data[ + -1][None, :].repeat(config.max_pos - 512, 1) + self.bert.model.embeddings.position_embeddings = my_pos_embeddings + self.vocab_size = self.bert.config.vocab_size + tgt_embeddings = nn.Embedding( + self.vocab_size, + self.bert.config.hidden_size, + padding_idx=1 if config.encoder == 'roberta' else 0) + + if config.share_emb: + tgt_embeddings.weight = copy.deepcopy( + self.bert.model.embeddings.word_embeddings.weight) + self.decoder = TransformerDecoder( + config.dec_layers, + config.dec_hidden_size, + heads=config.dec_heads, + d_ff=config.dec_ff_size, + dropout=config.dec_dropout, + embeddings=tgt_embeddings) + self.generator = PlugPointerGenerator(config.dec_hidden_size, + self.vocab_size) + self.generator.dense.weight = self.decoder.embeddings.weight + + if checkpoint is not None: + for key in list(checkpoint['model'].keys()): + if key.startswith('module.'): + checkpoint['model'][key.replace( + 'module.', '')] = checkpoint['model'][key] + checkpoint['model'].pop(key) + if key.startswith('plug.'): + checkpoint['model'][key.replace( + 'plug.', '')] = checkpoint['model'][key] + checkpoint['model'].pop(key) + msg = self.load_state_dict(checkpoint['model'], strict=False) + print(msg) + else: + for module in self.decoder.modules(): + if isinstance(module, (nn.Linear, nn.Embedding)): + module.weight.data.normal_(mean=0.0, std=0.02) + elif isinstance(module, nn.LayerNorm): + module.bias.data.zero_() + module.weight.data.fill_(1.0) + if isinstance(module, nn.Linear) and module.bias is not None: + module.bias.data.zero_() + for p in self.generator.parameters(): + if p.dim() > 1: + xavier_uniform_(p) + else: + p.data.zero_() + if config.use_bert_emb: + if config.encoder == 'roberta': + tgt_embeddings = nn.Embedding( + self.vocab_size, + self.bert.config.hidden_size, + padding_idx=1) + else: + tgt_embeddings = nn.Embedding( + self.vocab_size, + self.bert.config.hidden_size, + padding_idx=0) + tgt_embeddings.weight = copy.deepcopy( + self.bert.embeddings.word_embeddings.weight) + self.decoder.embeddings = tgt_embeddings + self.generator.dense.weight = self.decoder.embeddings.weight + + def forward(self, src, tgt, mask_src, token_type_ids): + top_vec, _ = self.bert( + src, mask_src, token_type_ids=token_type_ids, return_dict=False) + state = TransformerDecoderState(src) + decoder_outputs, attns, _ = self.decoder(state, tgt[:, :-1], top_vec) + return decoder_outputs, attns[-1], top_vec + + +class LabelSmoothingLoss(nn.Module): + """ + With label smoothing, + KL-divergence between q_{smoothed ground truth prob.}(w) + and p_{prob. computed by model}(w) is minimized. + """ + + def __init__(self, label_smoothing, tgt_vocab_size, ignore_index=-100): + assert 0.0 < label_smoothing <= 1.0 + self.padding_idx = ignore_index + super(LabelSmoothingLoss, self).__init__() + + smoothing_value = label_smoothing / (tgt_vocab_size - 2) + one_hot = torch.full((tgt_vocab_size, ), smoothing_value) + one_hot[self.padding_idx] = 0 + self.register_buffer('one_hot', one_hot.unsqueeze(0)) + self.confidence = 1.0 - label_smoothing + + def forward(self, output, target): + """ + output (FloatTensor): batch_size x n_classes + target (LongTensor): batch_size + """ + model_prob = self.one_hot.repeat(target.size(0), 1) + model_prob.scatter_(1, target.unsqueeze(1), self.confidence) + model_prob.masked_fill_((target == self.padding_idx).unsqueeze(1), 0) + + return F.kl_div(output, model_prob, reduction='sum') + + +class NMTLossCompute(nn.Module): + """ + Standard NMT Loss Computation. + """ + + def __init__(self, generator, symbols, vocab_size, label_smoothing=0.0): + super().__init__() + self.generator = generator + self.padding_idx = symbols['PAD'] + if label_smoothing > 0: + self.criterion = LabelSmoothingLoss( + label_smoothing, vocab_size, ignore_index=self.padding_idx) + else: + self.criterion = nn.NLLLoss( + ignore_index=self.padding_idx, reduction='sum') + + def _bottle(self, _v): + return _v.view(-1, _v.size(2)) + + def _unbottle(self, _v, batch_size): + return _v.view(-1, batch_size, _v.size(1)) + + def forward(self, tgt, output): + target = tgt[:, 1:] + batch_size, decoder_length = target.size(0), target.size(1) + normalization = target.ne(self.padding_idx).sum() + bottled_output = self._bottle(output) + scores = self.generator(bottled_output) + gtruth = target.contiguous().view(-1) + loss = self.criterion(scores, gtruth) + loss = loss.div(float(normalization)) + return loss, scores.view(batch_size, decoder_length, -1) + + +class PlugForConditionalGeneration(PlugPreTrainedModel): + + @dataclass + class Batch: + batch_size: int + src: torch.Tensor + tgt: torch.Tensor + mask_src: torch.Tensor + token_type_ids: torch.Tensor + query_id: List[None] = None + src_str: List[List[str]] = None + tgt_str: List[str] = None + + def __init__(self, config, checkpoint=None, dataset: str = 'default'): + super().__init__(config) + self.logger = logging.get_logger(__name__) + self.config = config + if config.encoder == 'roberta': + tokenizer = RobertaTokenizer.from_pretrained( + config.encoder_pth, do_lower_case=False) + symbols = { + 'BOS': tokenizer.cls_token_id, + 'EOS': tokenizer.sep_token_id, + 'PAD': tokenizer.pad_token_id, + 'EOQ': tokenizer.unk_token_id + } + elif config.encoder == 'bert' or config.encoder == 'zh_bert': + tokenizer = BertTokenizer.from_pretrained( + config.encoder_pth, do_lower_case=True) + symbols = { + 'BOS': tokenizer.vocab['[CLS]'], + 'EOS': tokenizer.vocab['[SEP]'], + 'PAD': tokenizer.vocab['[PAD]'], + 'EOQ': tokenizer.vocab['[unused2]'] + } + self.tokenizer = tokenizer + self.symbols = symbols + self.plug = PlugModel(config, checkpoint) + self.loss = NMTLossCompute(self.plug.generator, symbols, + self.plug.vocab_size, + config.label_smoothing) + # for generation + self.config.dataset = dataset + self.start_token = self.symbols['BOS'] + self.end_token = self.symbols['EOS'] + + def forward(self, src, tgt, mask_src=None, token_type_ids=None): + if mask_src is None: + mask_src = src.ne(self.symbols['PAD']).long() + output = self.plug(src, tgt, mask_src, token_type_ids)[0] + loss = self.loss(tgt, output) + return loss + + def translate_batch(self, + batch: 'Batch', + fast: bool = False, + *args, + **kwargs): + """ + Translate a batch of sentences. + + Mostly a wrapper around :obj:`Beam`. + + Args: + batch (:obj:`Batch`): a batch from a dataset object + data (:obj:`Dataset`): the dataset object + fast (bool): enables fast beam search (may not support all features) + + Todo: + Shouldn't need the original dataset. + """ + self.plug.eval() + with torch.no_grad(): + return self._fast_translate_batch(batch, *args, **kwargs) + + def _tile(self, x, count, dim=0): + perm = list(range(len(x.size()))) + if dim != 0: + perm[0], perm[dim] = perm[dim], perm[0] + x = x.permute(perm).contiguous() + out_size = list(x.size()) + out_size[0] *= count + batch = x.size(0) + x = x.view(batch, -1) \ + .transpose(0, 1) \ + .repeat(count, 1) \ + .transpose(0, 1) \ + .contiguous() \ + .view(*out_size) + if dim != 0: + x = x.permute(perm).contiguous() + return x + + def _top_k_top_p_filtering(self, + logits, + top_k=10, + top_p=1.0, + filter_value=-float('Inf'), + min_tokens_to_keep=1): + if top_k > 0: + top_k = min(max(top_k, min_tokens_to_keep), + logits.size(-1)) # Safety check + # Remove all tokens with a probability less than the last token of the top-k + indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, + None] + logits[indices_to_remove] = filter_value + + if top_p < 1.0: + sorted_logits, sorted_indices = torch.sort(logits, descending=True) + cumulative_probs = torch.cumsum( + F.softmax(sorted_logits, dim=-1), dim=-1) + + # Remove tokens with cumulative probability above the threshold (token with 0 are kept) + sorted_indices_to_remove = cumulative_probs > top_p + if min_tokens_to_keep > 1: + # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below) + sorted_indices_to_remove[..., :min_tokens_to_keep] = 0 + # Shift the indices to the right to keep also the first token above the threshold + sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[ + ..., :-1].clone() + sorted_indices_to_remove[..., 0] = 0 + + # scatter sorted tensors to original indexing + indices_to_remove = sorted_indices_to_remove.scatter( + 1, sorted_indices, sorted_indices_to_remove) + logits[indices_to_remove] = filter_value + return logits + + def _fast_translate_batch(self, + batch: 'Batch', + max_length: int = 80, + min_length: int = 10, + bad_words_ids=None, + early_stopping=True, + num_beams=3, + length_penalty=1.2, + repetition_penalty=1.2, + no_repeat_ngram_size=4, + *args, + **kwargs): + # TODO: faster code path for beam_size == 1. + # TODO: support these blacklisted features. + + num_beams = num_beams + batch_size = batch.batch_size + src = batch.src + mask_src = batch.mask_src + token_type_ids = batch.token_type_ids + + src_features, _ = self.plug.bert( + src, mask_src, token_type_ids=token_type_ids, return_dict=False) + state = TransformerDecoderState(src, self.plug.decoder.num_layers) + device = src_features.device + + # Tile states and memory beam_size times. + state.map_batch_fn( + lambda state, dim: self._tile(state, num_beams, dim=dim)) + src_features = self._tile(src_features, num_beams, dim=0) + batch_offset = torch.arange( + batch_size, dtype=torch.long, device=device) + beam_offset = torch.arange( + 0, + batch_size * num_beams, + step=num_beams, + dtype=torch.long, + device=device) + alive_seq = torch.full([batch_size * num_beams, 1], + self.start_token, + dtype=torch.long, + device=device) + + # cal bad_words_ids pre dict + bad_words_prefix_dict = {} + bad_words_prefix_len = set([]) + if bad_words_ids is not None: + for bw_id in bad_words_ids: + key = tuple(bw_id[:-1]) + value = bw_id[-1] + bad_words_prefix_dict[key] = bad_words_prefix_dict.get( + key, []) + [value] + bad_words_prefix_len.add(len(key)) + + # Give full probability to the first beam on the first step. + topk_log_probs = ( + torch.tensor( + [0.0] + [float('-inf')] * (num_beams - 1), + device=device).repeat(batch_size)) + + # Structure that holds finished hypotheses. + hypotheses = [[] for _ in range(batch_size)] # noqa: F812 + + results = {} + results['predictions'] = [[] for _ in range(batch_size)] # noqa: F812 + results['scores'] = [[] for _ in range(batch_size)] # noqa: F812 + results['gold_score'] = [0] * batch_size + results['batch'] = batch + + for step in range(max_length): + self.logger.info(f'step: {step + 1} / {max_length}') + decoder_input = alive_seq[:, -1].view(1, -1) + + # Decoder forward. + decoder_input = decoder_input.transpose(0, 1) + dec_out, attns, state = self.plug.decoder( + state, decoder_input, src_features, step=step) + + # Generator forward. + log_probs = self.plug.generator.forward( + dec_out.transpose(0, 1).squeeze(0)) + vocab_size = log_probs.size(-1) + + if step < min_length: + log_probs[:, self.end_token] = -1e20 + + # filter bad word + if len(bad_words_prefix_dict) > 0: + # cal bad word banned token: batch_size * num_beams + num_hypos = alive_seq.size(0) + bad_word_banned_token = [] + for i in range(num_hypos): + curr_banned_token = [] + for pre_len in bad_words_prefix_len: + pre_key = tuple(alive_seq[i, step + 1 - pre_len:step + + 1].cpu().numpy().tolist()) + curr_banned_token += bad_words_prefix_dict.get( + pre_key, []) + bad_word_banned_token.append(set(curr_banned_token)) + # set banned word prob=-1e20 + assert log_probs.size(0) == num_hypos + for i in range(num_hypos): + for banned_token in bad_word_banned_token[i]: + log_probs[i, banned_token] = -1e20 + + # do repetition_penalty + if repetition_penalty > 1.0: + """repetition penalty (from CTRL paper https://arxiv.org/abs/1909.05858). """ + # calculate prev_output_tokens for repetition_penalty: batch_size * num_beams + prev_output_tokens = self.calc_banned_tokens( + alive_seq, alive_seq.size(0), no_repeat_ngram_size, + step + 1) + # batch_size * num_beams + for i in range(log_probs.size(0)): + for previous_token in set(prev_output_tokens[i]): + if log_probs[i, previous_token] < 0: + log_probs[i, previous_token] *= repetition_penalty + else: + log_probs[i, previous_token] /= repetition_penalty + + # Multiply probs by the beam probability. + + curr_length_penalty = (step + 1)**length_penalty + # ''' + if self.config.sample_topk: + temperature = self.config.temperature + _scores = log_probs / temperature + _scores = self._top_k_top_p_filtering( + _scores, + top_k=self.config.top_k, + top_p=self.config.top_p, + min_tokens_to_keep=1 + ) # (batch_size * num_beams, vocab_size) + # Sample 2 next words for each beam (so we have some spare tokens + # and match output of greedy beam search) + topk_ids = torch.multinomial( + F.softmax(_scores, dim=-1), + num_samples=1) # (batch_size * num_beams, 2) + # Compute next scores + _scores = F.log_softmax( + _scores, dim=1) # (batch_size * num_beams, vocab_size) + + _scores += topk_log_probs.view(-1).unsqueeze(1) + _scores = _scores / curr_length_penalty + topk_scores = torch.gather( + _scores, -1, topk_ids) # (batch_size * num_beams, 2) + # log_probs += # (batch_size * num_beams, 2) + # Match shape of greedy beam search + topk_ids = topk_ids.view( + -1, num_beams) # (batch_size, 2 * num_beams) + topk_scores = topk_scores.view( + -1, num_beams) # (batch_size, 2 * num_beams) + # ''' + else: + log_probs += topk_log_probs.view(-1).unsqueeze(1) + curr_scores = log_probs / curr_length_penalty + + curr_scores = curr_scores.reshape(-1, num_beams * vocab_size) + topk_scores, topk_ids = curr_scores.topk(num_beams, dim=-1) + if (self.config.block_trigram): + cur_len = alive_seq.size(1) + if (cur_len > 3): + for i in range(alive_seq.size(0)): + fail = False + words = [int(w) for w in alive_seq[i]] + if self.config.encoder == 'roberta': + # words = [self.vocab.convert_ids_to_tokens[w] for w in words] + words = self.tokenizer.decode( + words).strip().split() + else: + words = [ + self.tokenizer.ids_to_tokens[w] for w in words + ] + words = ' '.join(words).replace(' ##', '').split() + if (len(words) <= 3): + continue + trigrams = [(words[i - 1], words[i], words[i + 1]) + for i in range(1, + len(words) - 1)] + trigram = tuple(trigrams[-1]) + if trigram in trigrams[:-1]: + fail = True + if fail: + curr_scores[i] = -10e20 + # Recover log probs. + topk_log_probs = topk_scores * curr_length_penalty + + # Resolve beam origin and true word ids. + # topk_beam_index = topk_ids.div(vocab_size) + topk_beam_index = topk_ids // vocab_size + topk_ids = topk_ids.fmod(vocab_size) + + # Map beam_index to batch_index in the flat representation. + batch_index = ( + topk_beam_index + + beam_offset[:topk_beam_index.size(0)].unsqueeze(1)) + select_indices = batch_index.view(-1) + + # Append last prediction. + alive_seq = torch.cat([ + alive_seq.index_select(0, select_indices), + topk_ids.view(-1, 1) + ], -1) + + is_finished = topk_ids.eq(self.end_token) + if step + 1 == max_length: + is_finished.fill_(self.end_token) + # End condition is top beam is finished. + end_condition = is_finished[:, 0].eq(1) + # Save finished hypotheses. + if is_finished.any(): + predictions = alive_seq.view(-1, num_beams, alive_seq.size(-1)) + for i in range(is_finished.size(0)): + b = batch_offset[i] + if end_condition[i]: + is_finished[i].fill_(self.end_token) + finished_hyp = is_finished[i].nonzero().view(-1) + # Store finished hypotheses for this batch. + for j in finished_hyp: + hypotheses[b].append( + (topk_scores[i, j], predictions[i, j, 1:])) + if early_stopping and len(hypotheses) == num_beams: + end_condition[i] = True + # If the batch reached the end, save the n_best hypotheses. + if end_condition[i]: + best_hyp = sorted( + hypotheses[b], key=lambda x: x[0], reverse=True) + if self.config.dataset == 'qg_ranking_test' or ( + self.config.dataset == 'paraphrase' + and not self.config.sample_topk): + for each in best_hyp[:num_beams]: + score, pred = each + results['scores'][b].append(score) + results['predictions'][b].append(pred) + else: + score, pred = best_hyp[0] + results['scores'][b].append(score) + results['predictions'][b].append(pred) + non_finished = end_condition.eq(0).nonzero().view(-1) + # If all sentences are translated, no need to go further. + if len(non_finished) == 0: + break + # Remove finished batches for the next step. + topk_log_probs = topk_log_probs.index_select(0, non_finished) + batch_index = batch_index.index_select(0, non_finished) + batch_offset = batch_offset.index_select(0, non_finished) + alive_seq = predictions.index_select(0, non_finished) \ + .view(-1, alive_seq.size(-1)) + + # Reorder states. + select_indices = batch_index.view(-1) + src_features = src_features.index_select(0, select_indices) + state.map_batch_fn( + lambda state, dim: state.index_select(dim, select_indices)) + + return results + + def calc_banned_tokens(self, prev_input_ids, num_hypos, + no_repeat_ngram_size, cur_len): + # Copied from fairseq for no_repeat_ngram in beam_search""" + if cur_len + 1 < no_repeat_ngram_size: + # return no banned tokens if we haven't generated no_repeat_ngram_size tokens yet + return [[] for _ in range(num_hypos)] + generated_ngrams = [{} for _ in range(num_hypos)] + for idx in range(num_hypos): + gen_tokens = prev_input_ids[idx].cpu().numpy().tolist() + generated_ngram = generated_ngrams[idx] + for ngram in zip( + *[gen_tokens[i:] for i in range(no_repeat_ngram_size)]): + prev_ngram_tuple = tuple(ngram[:-1]) + generated_ngram[prev_ngram_tuple] = generated_ngram.get( + prev_ngram_tuple, []) + [ngram[-1]] + + def _get_generated_ngrams(hypo_idx): + # Before decoding the next token, prevent decoding of ngrams that have already appeared + start_idx = cur_len + 1 - no_repeat_ngram_size + ngram_idx = tuple( + prev_input_ids[hypo_idx, + start_idx:cur_len].cpu().numpy().tolist()) + return generated_ngrams[hypo_idx].get(ngram_idx, []) + + banned_tokens = [ + _get_generated_ngrams(hypo_idx) for hypo_idx in range(num_hypos) + ] + return banned_tokens + + def translate(self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor = None, + token_type_ids=None, + *args, + **kwargs) -> Dict[str, torch.Tensor]: + if attention_mask is None: + attention_mask = input_ids.ne(self.symbols['PAD']).long() + batch = self.Batch( + batch_size=input_ids.size()[0], + src=input_ids, + tgt=None, + token_type_ids=token_type_ids, + mask_src=attention_mask) + translation_batch = self.translate_batch(batch, *args, **kwargs) + + preds = translation_batch['predictions'] + return {'predictions': preds} diff --git a/modelscope/models/nlp/fid_plug/configuration.py b/modelscope/models/nlp/fid_plug/configuration.py new file mode 100644 index 00000000..ec8e0635 --- /dev/null +++ b/modelscope/models/nlp/fid_plug/configuration.py @@ -0,0 +1,103 @@ +# Copyright 2021-2022 The Alibaba DAMO NLP Team Authors. +# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team. +# Copyright (c) 2018, NVIDIA CORPORATION. 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. +""" PLUG model configuration """ +from transformers.configuration_utils import PretrainedConfig + + +class PlugConfig(PretrainedConfig): + r""" + Configuration objects inherit from :class:`~transformers.PretrainedConfig` and can be used to control the model + outputs. Read the documentation from :class:`~transformers.PretrainedConfig` for more information. + + + Args: + vocab_size (:obj:`int`, `optional`, defaults to 30522): + Vocabulary size of the BERT model. Defines the number of different tokens that can be represented by the + :obj:`inputs_ids` passed when calling :class:`~transformers.BertModel` or + :class:`~transformers.TFBertModel`. + hidden_size (:obj:`int`, `optional`, defaults to 768): + Dimensionality of the encoder layers and the pooler layer. + num_hidden_layers (:obj:`int`, `optional`, defaults to 12): + Number of hidden layers in the Transformer encoder. + num_attention_heads (:obj:`int`, `optional`, defaults to 12): + Number of attention heads for each attention layer in the Transformer encoder. + intermediate_size (:obj:`int`, `optional`, defaults to 3072): + Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder. + hidden_act (:obj:`str` or :obj:`Callable`, `optional`, defaults to :obj:`"gelu"`): + The non-linear activation function (function or string) in the encoder and pooler. If string, + :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported. + hidden_dropout_prob (:obj:`float`, `optional`, defaults to 0.1): + The dropout probability for all fully connected layers in the embeddings, encoder, and pooler. + attention_probs_dropout_prob (:obj:`float`, `optional`, defaults to 0.1): + The dropout ratio for the attention probabilities. + max_position_embeddings (:obj:`int`, `optional`, defaults to 512): + The maximum sequence length that this model might ever be used with. Typically set this to something large + just in case (e.g., 512 or 1024 or 2048). + type_vocab_size (:obj:`int`, `optional`, defaults to 2): + The vocabulary size of the :obj:`token_type_ids` passed when calling :class:`~transformers.BertModel` or + :class:`~transformers.TFBertModel`. + initializer_range (:obj:`float`, `optional`, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + layernorm_epsilon (:obj:`float`, `optional`, defaults to 1e-12): + The epsilon used by the layer normalization layers. + dec_hidden_layers (:obj:`int`, `optional`, defaults to 12): + Number of hidden layers in the Transformer decoder. + attn_separate (:obj:`bool`, `optional`, defaults to false): + Whether or not to separate the q, k, v of attention. + + Examples:: + + >>> import PlugModel, PlugConfig + >>> configuration = PlugConfig() + + >>> # Initializing a model from the configuration + >>> model = PlugModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + """ + model_type = 'plug' + + def __init__(self, + encoder='roberta', + encoder_pth='roberta-base', + max_pos=512, + share_emb=False, + dec_layers=12, + dec_hidden_size=768, + dec_heads=8, + dec_ff_size=3072, + dec_dropout=0.2, + use_bert_emb=True, + label_smoothing=0.1, + sample_topk=False, + block_trigram=False, + **kwargs): + super().__init__(**kwargs) + self.encoder = encoder + self.encoder_pth = encoder_pth + self.max_pos = max_pos + self.share_emb = share_emb + self.dec_layers = dec_layers + self.dec_hidden_size = dec_hidden_size + self.dec_heads = dec_heads + self.dec_ff_size = dec_ff_size + self.dec_dropout = dec_dropout + self.use_bert_emb = use_bert_emb + self.label_smoothing = label_smoothing + # Translator + self.sample_topk = sample_topk + self.block_trigram = block_trigram diff --git a/modelscope/models/nlp/fid_plug/text_generation.py b/modelscope/models/nlp/fid_plug/text_generation.py new file mode 100644 index 00000000..2fe7cb72 --- /dev/null +++ b/modelscope/models/nlp/fid_plug/text_generation.py @@ -0,0 +1,180 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import io +import os + +import torch +from transformers.modeling_outputs import Seq2SeqLMOutput + +from modelscope.metainfo import Models +from modelscope.models import Model +from modelscope.models.base import TorchModel +from modelscope.models.builder import MODELS +from modelscope.outputs import TextGenerationModelOutput, TokenGeneratorOutput +from modelscope.utils import logger as logging +from modelscope.utils.constant import Tasks +from .backbone import PlugForConditionalGeneration +from .configuration import PlugConfig + +CONFIG_NAME = 'config.json' +WEIGHTS_NAME = 'pytorch_model.bin' + + +class PlugV2Chat(TorchModel): + + def __init__(self, model_dir, *args, **kwargs): + super().__init__(model_dir, *args, **kwargs) + # init model + plug_config_file = os.path.join(model_dir, CONFIG_NAME) + plug_config = PlugConfig.from_json_file(plug_config_file) + self.backbone = PlugForConditionalGeneration(plug_config) + # load weights + pretrained_model_path = os.path.join(model_dir, WEIGHTS_NAME) + with io.open(pretrained_model_path, 'rb') as f: + checkpoint = torch.load(f, map_location='cpu') + if 'model' in checkpoint: + checkpoint = checkpoint['model'] + for key in list(checkpoint.keys()): + # for old plugv2 version + if key.startswith('translator'): + checkpoint.pop(key) + continue + if key.startswith('module.'): + checkpoint[key.replace('module.', '')] = checkpoint[key] + checkpoint.pop(key) + if key.startswith('backbone.plug.bert.bert.'): + checkpoint[key.replace('backbone.plug.bert.bert.', + 'bert.')] = checkpoint[key] + checkpoint.pop(key) + elif key.startswith('backbone.plug.'): + checkpoint[key.replace('backbone.plug.', + '')] = checkpoint[key] + checkpoint.pop(key) + msg = self.backbone.plug.load_state_dict(checkpoint, strict=False) + print(f'| {msg}') + + def generate(self, input_ids, token_type_ids=None, *args, **kwargs): + pred_result = self.backbone.translate( + input_ids=input_ids, + token_type_ids=token_type_ids, + *args, + **kwargs)['predictions'] + response = [x[0].tolist() for x in pred_result] + response = torch.tensor(response) + return response + + def forward(self, + input_ids, + decoder_input_ids, + token_type_ids=None, + *args, + **kwargs): + loss = self.backbone.forward( + src=input_ids, + tgt=decoder_input_ids, + token_type_ids=token_type_ids, + **kwargs) + return Seq2SeqLMOutput(loss=loss[0], logits=loss[1]) + + +class PlugV2EncoderWrapper(torch.nn.Module): + + def __init__(self, bert): + super().__init__() + + self.bert = bert + self.n_passages = None + + def set_n_passages(self, n_passages): + self.n_passages = n_passages + + def forward(self, + input_ids=None, + attention_mask=None, + token_type_ids=None, + *args, + **kwargs): + # total_length = n_passages * passage_length + bsz, total_length = input_ids.shape + passage_length = total_length // self.n_passages + input_ids = input_ids.view(bsz * self.n_passages, passage_length) + if token_type_ids is not None: + token_type_ids = token_type_ids.view(bsz * self.n_passages, + passage_length) + if attention_mask is not None: + attention_mask = attention_mask.view(bsz * self.n_passages, + passage_length) + outputs = self.bert( + input_ids, + attention_mask, + token_type_ids=token_type_ids, + *args, + **kwargs) + if isinstance(outputs, tuple): + outputs = (outputs[0].view(bsz, self.n_passages * passage_length, + -1), ) + outputs[1:] + else: + outputs.last_hidden_state = outputs.last_hidden_state.view( + bsz, self.n_passages * passage_length, -1) + return outputs + + +@MODELS.register_module(Tasks.fid_dialogue, module_name=Models.fid_plug) +class PlugV2FidChat(PlugV2Chat): + + def __init__(self, model_dir, *args, **kwargs): + super().__init__(model_dir, *args, **kwargs) + self.wrap_encoder() + + def wrap_encoder(self): + self.backbone.plug.bert = PlugV2EncoderWrapper(self.backbone.plug.bert) + + def unwrap_encoder(self): + self.backbone.plug.bert = self.backbone.plug.bert.bert + + def load(self, + pretrained_model_path, + from_tf=False): # only invoked when model is not onnx format + self.unwrap_encoder() + super().load(pretrained_model_path) + self.wrap_encoder() + + def generate(self, inputs, *args, **kwargs): + input_ids = inputs.get('input_ids') + attention_mask = inputs.get('attention_mask', None) + token_type_ids = inputs.get('token_type_ids', None) + n_passages = input_ids.size(1) + self.backbone.plug.bert.set_n_passages(n_passages) + input_ids = input_ids.view(input_ids.size(0), -1) + if token_type_ids is not None: + token_type_ids = token_type_ids.view(token_type_ids.size(0), -1) + response = super().generate( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + *args, + **kwargs) + return TokenGeneratorOutput(sequences=response) + + def forward(self, + input_ids, + decoder_input_ids, + token_type_ids=None, + *args, + **kwargs): + if input_ids is not None: + # inputs might have already be resized in the generate method + if input_ids.dim() == 3: + n_passages = input_ids.size(1) + self.backbone.plug.bert.set_n_passages(n_passages) + input_ids = input_ids.view(input_ids.size(0), -1) + if token_type_ids is not None: + token_type_ids = token_type_ids.view(input_ids.size(0), -1) + seq2seq_lm_output = super().forward( + input_ids, + decoder_input_ids=decoder_input_ids, + token_type_ids=token_type_ids, + *args, + **kwargs) + return TextGenerationModelOutput( + loss=seq2seq_lm_output.loss, logits=seq2seq_lm_output.logits) diff --git a/modelscope/outputs/outputs.py b/modelscope/outputs/outputs.py index 94023c37..cfed6358 100644 --- a/modelscope/outputs/outputs.py +++ b/modelscope/outputs/outputs.py @@ -674,6 +674,12 @@ TASK_OUTPUTS = { # } Tasks.text_generation: [OutputKeys.TEXT], + # fid dialogue result for single sample + # { + # "text": "My name is Mike" + # } + Tasks.fid_dialogue: [OutputKeys.TEXT], + # summarization result for single sample # { # "text": "this is the text generated by a model." diff --git a/modelscope/pipeline_inputs.py b/modelscope/pipeline_inputs.py index acfcd8c3..1388846a 100644 --- a/modelscope/pipeline_inputs.py +++ b/modelscope/pipeline_inputs.py @@ -190,6 +190,12 @@ TASK_INPUTS = { Tasks.text_ranking: (InputType.TEXT, InputType.TEXT), Tasks.text_generation: InputType.TEXT, + Tasks.fid_dialogue: { + 'history': InputType.TEXT, + 'knowledge': InputType.TEXT, + 'bot_profile': InputType.TEXT, + 'user_profile': InputType.TEXT, + }, Tasks.fill_mask: InputType.TEXT, Tasks.task_oriented_conversation: { diff --git a/modelscope/pipelines/nlp/__init__.py b/modelscope/pipelines/nlp/__init__.py index 8fcf3e3f..adfb90c5 100644 --- a/modelscope/pipelines/nlp/__init__.py +++ b/modelscope/pipelines/nlp/__init__.py @@ -25,6 +25,7 @@ if TYPE_CHECKING: from .translation_quality_estimation_pipeline import TranslationQualityEstimationPipeline from .text_error_correction_pipeline import TextErrorCorrectionPipeline from .text_generation_pipeline import TextGenerationPipeline, TextGenerationT5Pipeline + from .fid_dialogue_pipeline import FidDialoguePipeline from .token_classification_pipeline import TokenClassificationPipeline from .translation_pipeline import TranslationPipeline from .word_segmentation_pipeline import WordSegmentationPipeline, WordSegmentationThaiPipeline @@ -65,6 +66,7 @@ else: 'text_error_correction_pipeline': ['TextErrorCorrectionPipeline'], 'text_generation_pipeline': ['TextGenerationPipeline', 'TextGenerationT5Pipeline'], + 'fid_dialogue_pipeline': ['FidDialoguePipeline'], 'text2text_generation_pipeline': ['Text2TextGenerationPipeline'], 'token_classification_pipeline': ['TokenClassificationPipeline'], 'translation_pipeline': ['TranslationPipeline'], diff --git a/modelscope/pipelines/nlp/fid_dialogue_pipeline.py b/modelscope/pipelines/nlp/fid_dialogue_pipeline.py new file mode 100644 index 00000000..2880b4a2 --- /dev/null +++ b/modelscope/pipelines/nlp/fid_dialogue_pipeline.py @@ -0,0 +1,176 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import re +from typing import Any, Dict, Optional, Union + +import torch + +from modelscope.metainfo import Pipelines +from modelscope.models.base import Model +from modelscope.outputs import OutputKeys, TokenGeneratorOutput +from modelscope.pipelines.base import Pipeline +from modelscope.pipelines.builder import PIPELINES +from modelscope.preprocessors import Preprocessor +from modelscope.utils.constant import ModelFile, Tasks + +context_template = '假设我和你正在进行对话，请你给我得体、准确、友好的回复。以下是我们的对话内容。{context}' +history_template = '假设我和你正在进行对话，请你给我得体、准确、友好的回复。以下是我们的对话内容。{context}' \ + '#以下是在此之前我们的对话内容，可作为回复时的参考。{history}' +knowledge_template = '假设我和你正在进行对话，请你给我得体、准确、友好的回复。以下是我们的对话内容。{context}' \ + '#以下是和对话相关的知识，请你参考该知识进行回复。{knowledge}' +user_profile_template = '假设我和你正在进行对话，请你给我得体、准确、友好的回复。以下是我们的对话内容。{context}' \ + '#假设以下是你对我所了解的信息，请你参考该信息并避免你的回复和该信息矛盾，信息如下：{user_profile}' +bot_profile_template = '假设我和你正在进行对话，请你给我得体、准确、友好的回复。以下是我们的对话内容。{context}' \ + '#假设以下是你的人物设定，请你参考该信息并避免你的回复和该信息矛盾，信息如下：{bot_profile}' + +__all__ = ['FidDialoguePipeline'] + + +@PIPELINES.register_module( + Tasks.fid_dialogue, module_name=Pipelines.fid_dialogue) +class FidDialoguePipeline(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 fid-dialogue pipeline for prediction. + + Args: + model (str or Model): Supply either a local model dir which supported the text generation 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 + >>> from modelscope.utils.constant import Tasks + >>> pipeline_ins = pipeline(Tasks.fid_dialogue, model='damo/plug-dialogue', model_revision='v1.0.1') + >>> input = { + >>> "history": "你好[SEP]你好，我是小达，很高兴认识你！[SEP]李白是谁", + >>> "bot_profile": "我是小达;我是女生;我是单身;我今年21岁;我生日是2001年11月11日", + >>> "knowledge": "唐代诗人李白（701年—762年12月）,字太白,号青莲居士,又号“谪仙人”[SEP]李白（公元701年—公元762年），字太白", + >>> "user_profile": "你是小明" + >>> } + >>> result = pipeline_ins(input) + >>> print(result) + """ + super().__init__( + model=model, + preprocessor=preprocessor, + config_file=config_file, + device=device, + auto_collate=auto_collate, + **kwargs) + + if preprocessor is None: + self.preprocessor_tokenizer = Preprocessor.from_pretrained( + self.model.model_dir, **kwargs) + + assert isinstance(self.model, Model), \ + f'please check whether model config exists in {ModelFile.CONFIGURATION}' + self.model = self.model.to(self.device) + self.model.eval() + + self.SEP = '[SEP]' + + def forward(self, inputs: Dict[str, Any], **forward_params): + with torch.no_grad(): + return self.model.generate(inputs, **forward_params) + + def preprocess(self, inputs: Dict[str, Any], + **preprocess_params) -> Dict[str, Any]: + # init params + max_encoder_length = 300 + if 'max_encoder_length' in preprocess_params: + max_encoder_length = preprocess_params.pop('max_encoder_length') + # get raw data + history = inputs['history'] if 'history' in inputs else '' + if len(history) <= 0: + raise Exception('history is necessary!') + knowledge = inputs['knowledge'] if 'knowledge' in inputs else '' + user_profile = inputs[ + 'user_profile'] if 'user_profile' in inputs else '' + bot_profile = inputs['bot_profile'] if 'bot_profile' in inputs else '' + # parse raw data + history = history.split(self.SEP) + context = history[-3:] + context = self.process_context(context) + history = history[:-3] + history = self.process_history(history) + knowledge = knowledge.split(self.SEP) + + model_input = [] + if history and len(history) > 0: + model_input.append( + history_template.format(context=context, history=history)) + if knowledge and len(knowledge) > 0: + for know in knowledge: + model_input.append( + knowledge_template.format(context=context, knowledge=know)) + if user_profile and len(user_profile) > 0: + model_input.append( + user_profile_template.format( + context=context, user_profile=user_profile)) + if bot_profile and len(bot_profile) > 0: + model_input.append( + bot_profile_template.format( + context=context, bot_profile=bot_profile)) + + if not model_input: + model_input.append(context_template.format(context=context)) + + for i in range(len(model_input)): + model_input[i] = re.sub('[ \t]+', '▂', model_input[i]) + + # tokenization + input_ids = self.preprocessor_tokenizer( + {'src_txt': model_input}, + padding=True, + truncation=True, + max_length=max_encoder_length, + return_tensors='pt')['input_ids'].unsqueeze(0).to(self.device) + input_dict = { + 'input_ids': + input_ids.to(torch.int64).to(self.device), + 'attention_mask': (input_ids != 0).to(torch.int64).to(self.device), + 'token_type_ids': + torch.zeros(input_ids.shape).to(torch.int64).to(self.device) + } + + return input_dict + + def process_context(self, context_list): + subject = '我' + for i in range(len(context_list) - 1, -1, -1): + if len(context_list[i]) > 0 and context_list[i][ + -1] not in '!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~、。，？！；：“”（）【】《》〈〉……': + context_list[i] = context_list[i] + '。' + context_list[i] = subject + '：' + context_list[i] + subject = '你' if subject == '我' else '我' + return ''.join(context_list) + + def process_history(self, history_list): + subject = '你' + for i in range(len(history_list) - 1, -1, -1): + if len(history_list[i]) > 0 and history_list[i][ + -1] not in '!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~、。，？！；：“”（）【】《》〈〉……': + history_list[i] = history_list[i] + '。' + history_list[i] = subject + '：' + history_list[i] + subject = '你' if subject == '我' else '我' + return ''.join(history_list) + + def postprocess(self, inputs: TokenGeneratorOutput, + **postprocess_params) -> Dict[str, Any]: + + if torch.cuda.is_available(): + hypotheses = inputs.sequences.detach().cpu().tolist() + + response = self.preprocessor_tokenizer.decode( + hypotheses[0], skip_special_tokens=True) + response = response.replace(' ', '') + return {OutputKeys.TEXT: response} diff --git a/modelscope/preprocessors/nlp/transformers_tokenizer.py b/modelscope/preprocessors/nlp/transformers_tokenizer.py index 1af7bfa7..7a4705b9 100644 --- a/modelscope/preprocessors/nlp/transformers_tokenizer.py +++ b/modelscope/preprocessors/nlp/transformers_tokenizer.py @@ -83,7 +83,7 @@ class NLPTokenizer: if model_type in (Models.structbert, Models.gpt3, Models.palm, Models.plug, Models.megatron_bert, - Models.plug_mental): + Models.plug_mental, Models.fid_plug): from transformers import BertTokenizer, BertTokenizerFast tokenizer = BertTokenizerFast if self.use_fast else BertTokenizer return tokenizer.from_pretrained( diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index 78e1f56e..9647abcd 100644 --- a/modelscope/utils/constant.py +++ b/modelscope/utils/constant.py @@ -155,6 +155,7 @@ class NLPTasks(object): token_classification = 'token-classification' conversational = 'conversational' text_generation = 'text-generation' + fid_dialogue = 'fid-dialogue' text2text_generation = 'text2text-generation' task_oriented_conversation = 'task-oriented-conversation' dialog_intent_prediction = 'dialog-intent-prediction' diff --git a/tests/pipelines/test_plug_dialogue.py b/tests/pipelines/test_plug_dialogue.py new file mode 100644 index 00000000..cde65bea --- /dev/null +++ b/tests/pipelines/test_plug_dialogue.py @@ -0,0 +1,45 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import unittest + +from modelscope.pipelines import pipeline +from modelscope.utils.constant import Tasks +from modelscope.utils.demo_utils import DemoCompatibilityCheck +from modelscope.utils.test_utils import test_level + + +class PlugDialogueTest(unittest.TestCase, DemoCompatibilityCheck): + + know_list = [ + '唐代诗人李白（701年—762年12月）,▂字太白,▂号青莲居士,▂又号“谪仙人”,▂唐代伟大的浪漫主义诗人,▂被后人誉为“诗仙”,▂与杜甫并称为“李杜”,▂为了与另两位诗人李商隐与杜牧即“小李杜”区别,▂杜甫与李白', + '白词”享有极为崇高的地位。李白▂主要成就▂创造了古代积极浪漫主义文学高峰、为唐诗的繁荣与发展打开了新局面、开创了中国古典诗歌的黄金时代', + '李白（701年—762年），字太白，号青莲居士，又号“谪仙人”。是唐代伟大的浪漫主义诗人，被后人誉为“诗仙”。与杜甫并称为“李杜”，为了与另两位诗人李商隐与杜牧即“小李杜”区别，杜甫与', + ] + input = { + 'history': '你好[SEP]你好，我是小达，很高兴认识你！[SEP]李白是谁', + 'knowledge': '[SEP]'.join(know_list), + 'bot_profile': + '我是小达;我是女生;我是单身;我今年21岁;我生日是2001年11月11日;我是天蝎座;我现在在复旦大学上学;我家现在常住上海', + 'user_profile': '你是小明' + } + + def setUp(self) -> None: + self.task = Tasks.fid_dialogue + self.model_id = 'damo/plug-dialogue' + self.model_revision = 'v1.0.1' + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_run_with_pipeline(self): + pipeline_ins = pipeline( + task=self.task, + model=self.model_id, + model_revision=self.model_revision) + result = pipeline_ins(self.input) + print(result) + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_demo_compatibility(self): + self.compatibility_check() + + +if __name__ == '__main__': + unittest.main()