[to #42322933] add fuse-in-decoder dialogue task

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'

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={},
+    )

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

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)

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."

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: {

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

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}

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(

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'

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()