Merge the speaker-turn-detection codes, local test finished

# Speaker Diarization Speaker-Turn Detection CR

和Dialogue-Detection一样，本模型是Speaker Diarization（`audio/speaker diarization`，语音/说话人日志）任务下的一个子模块。

本次提交的是基于文本进行判断的模型，本地模型的初始模型基于huggingface训练的，此提交中复用了部分 `nlp/token-classification` 模型的代码。为了方便后续维护以及与nlp方面代码解耦，在model、pipeline以及preprocessor中 **单独** 创建了相应模块并重新register。
Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/13364720
* std first commit

* local test pass for speaker-turn-detection

* update speaker-turn-detection pipeline task outputs format; update pipeline outputs; update test scripts

modelscope/metainfo.py

@@ -498,6 +498,7 @@ class Pipelines(object):
     speaker_verification_eres2net = 'speaker-verification-eres2net'
     speaker_change_locating = 'speaker-change-locating'
     speaker_diarization_dialogue_detection = 'speaker-diarization-dialogue-detection'
+    speaker_diarization_semantic_speaker_turn_detection = 'speaker-diarization-semantic-speaker-turn-detection'
     segmentation_clustering = 'segmentation-clustering'
     lm_inference = 'language-score-prediction'
     speech_timestamp_inference = 'speech-timestamp-inference'

modelscope/models/audio/sv/speaker_diarization_semantic_speaker_turn_detection.py

@@ -0,0 +1,137 @@
+# copyright (c) Alibaba, Inc. and its affiliates.
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn import CrossEntropyLoss
+
+from modelscope.metainfo import Heads, Models, TaskModels
+from modelscope.models.base import TorchHead
+from modelscope.models.builder import HEADS, MODELS
+from modelscope.models.nlp.task_models.task_model import EncoderModel
+from modelscope.outputs import (AttentionTokenClassificationModelOutput,
+                                ModelOutputBase, OutputKeys)
+from modelscope.utils import logger as logging
+from modelscope.utils.constant import Tasks
+from modelscope.utils.hub import parse_label_mapping
+
+logger = logging.get_logger()
+
+
+@HEADS.register_module(
+    Tasks.speaker_diarization_semantic_speaker_turn_detection,
+    module_name=Heads.token_classification)
+class TokenClassificationHead(TorchHead):
+
+    def __init__(self,
+                 hidden_size=768,
+                 classifier_dropout=0.1,
+                 num_labels=None,
+                 **kwargs):
+        super().__init__(
+            num_labels=num_labels,
+            classifier_dropout=classifier_dropout,
+            hidden_size=hidden_size,
+        )
+        assert num_labels is not None
+        self.dropout = nn.Dropout(classifier_dropout)
+        self.classifier = nn.Linear(hidden_size, num_labels)
+
+    def forward(self,
+                inputs: ModelOutputBase,
+                attention_mask=None,
+                labels=None,
+                **kwargs):
+        sequence_output = inputs.last_hidden_state
+        sequence_output = self.dropout(sequence_output)
+        logits = self.classifier(sequence_output)
+        loss = None
+        if labels is not None:
+            loss = self.compute_loss(logits, attention_mask, labels)
+
+        return AttentionTokenClassificationModelOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=inputs.hidden_states,
+            attentions=inputs.attentions)
+
+    def compute_loss(self, logits: torch.Tensor, attention_mask,
+                     labels) -> torch.Tensor:
+        loss_fct = CrossEntropyLoss()
+        # Only keep active parts of the loss
+        if attention_mask is not None:
+            active_loss = attention_mask.view(-1) == 1
+            active_logits = logits.view(-1, self.num_labels)
+            active_labels = torch.where(
+                active_loss, labels.view(-1),
+                torch.tensor(loss_fct.ignore_index).type_as(labels))
+            loss = loss_fct(active_logits, active_labels)
+        else:
+            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+        return loss
+
+
+@MODELS.register_module(
+    Tasks.speaker_diarization_semantic_speaker_turn_detection,
+    module_name=TaskModels.token_classification)
+class ModelForTokenClassification(EncoderModel):
+    task = Tasks.token_classification
+    head_type = Heads.token_classification
+    base_model_prefix = 'bert'
+    override_base_model_prefix = True
+
+    def __init__(self, model_dir: str, *args, **kwargs):
+        self.id2label = {}
+
+        num_labels = kwargs.get('num_labels')
+        if num_labels is None:
+            label2id = parse_label_mapping(model_dir)
+            if label2id is not None and len(label2id) > 0:
+                num_labels = len(label2id)
+            self.id2label = {id: label for label, id in label2id.items()}
+        kwargs['num_labels'] = num_labels
+        super(ModelForTokenClassification,
+              self).__init__(model_dir, *args, **kwargs)
+
+    def parse_head_cfg(self):
+        head_cfg = super().parse_head_cfg()
+        if hasattr(head_cfg, 'classifier_dropout'):
+            head_cfg['classifier_dropout'] = (
+                head_cfg.classifier_dropout if head_cfg['classifier_dropout']
+                is not None else head_cfg.hidden_dropout_prob)
+        else:
+            head_cfg['classifier_dropout'] = head_cfg.hidden_dropout_prob
+        head_cfg['num_labels'] = self.config.num_labels
+        return head_cfg
+
+    def forward(self,
+                input_ids=None,
+                attention_mask=None,
+                token_type_ids=None,
+                position_ids=None,
+                head_mask=None,
+                inputs_embeds=None,
+                labels=None,
+                output_attentions=None,
+                output_hidden_states=None,
+                return_dict=None,
+                offset_mapping=None,
+                label_mask=None,
+                **kwargs):
+        kwargs['offset_mapping'] = offset_mapping
+        kwargs['label_mask'] = label_mask
+        outputs = super().forward(input_ids, attention_mask, token_type_ids,
+                                  position_ids, head_mask, inputs_embeds,
+                                  labels, output_attentions,
+                                  output_hidden_states, **kwargs)
+
+        outputs.offset_mapping = offset_mapping
+        outputs.label_mask = label_mask
+
+        return outputs
+
+
+@MODELS.register_module(
+    Tasks.speaker_diarization_semantic_speaker_turn_detection,
+    module_name=Models.bert)
+class BertForTokenClassification(ModelForTokenClassification):
+    base_model_type = 'bert'

modelscope/outputs/outputs.py

@@ -1245,6 +1245,16 @@ TASK_OUTPUTS = {
     # { "text": "你好，明天！"}
     Tasks.punctuation: [OutputKeys.TEXT],
 
+    # speaker diarization semantic speaker-turn detection
+    # {
+    #    "logits": [[0.7, 0.3], ..., [0.88, 0.12]],
+    #    "text": "您好。您好，初次见面请多指教。",
+    #    "prediction": [-100, -100, -100, 1, -100,..., -100, 0]
+    # }
+    Tasks.speaker_diarization_semantic_speaker_turn_detection: [
+        OutputKeys.LOGITS, OutputKeys.TEXT, OutputKeys.PREDICTION
+    ],
+
     # language model result for single sample
     # { "text": " hel@@ lo 大 家 好 呀 </s>
     #               p( hel@@ | <s> ) = 0.00057767 [ -7.45650959 ]

modelscope/pipeline_inputs.py

@@ -332,10 +332,12 @@ TASK_INPUTS = {
     InputType.TEXT,
     Tasks.keyword_spotting:
     InputType.AUDIO,
-    Tasks.inverse_text_processing:
-    InputType.TEXT,
     Tasks.speaker_diarization_dialogue_detection:
     InputType.TEXT,
+    Tasks.speaker_diarization_semantic_speaker_turn_detection:
+    InputType.TEXT,
+    Tasks.inverse_text_processing:
+    InputType.TEXT,
 
     # ============ multi-modal tasks ===================
     Tasks.image_captioning: [InputType.IMAGE, {

modelscope/pipelines/audio/speaker_diarization_semantic_speaker_turn_detection_pipeline.py

@@ -0,0 +1,116 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+
+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 Preprocessor
+from modelscope.utils.constant import ModelFile, Tasks
+from modelscope.utils.logger import get_logger
+from modelscope.utils.tensor_utils import (torch_nested_detach,
+                                           torch_nested_numpify)
+
+logger = get_logger()
+
+
+@PIPELINES.register_module(
+    Tasks.speaker_diarization_semantic_speaker_turn_detection,
+    module_name=Pipelines.speaker_diarization_semantic_speaker_turn_detection)
+class SpeakerDiarizationSemanticSpeakerTurnDetectionPipeline(Pipeline):
+    r"""The inference pipeline for Speaker Diarization Semantic Speaker-Turn Detection Task.
+
+        Examples:
+            >>> from modelscope.pipelines import pipeline
+            >>> pipeline_ins = pipeline("speaker_diarization_semantic_speaker_turn_detection",
+                            model="damo/speech_bert_semantic-spk-turn-detection-punc_speaker-diarization_chinese")
+            >>> input_text = ""
+            >>> print(pipeline_ins(input_text))
+    """
+    PUNC_LIST = ['。', '，', '？', '！']
+
+    def __init__(self,
+                 model: Union[Model, str],
+                 preprocessor: Optional[Preprocessor] = None,
+                 config_file: str = None,
+                 device: str = 'gpu',
+                 auto_collate=True,
+                 sequence_length=128,
+                 **kwargs):
+        super().__init__(
+            model=model,
+            preprocessor=preprocessor,
+            config_file=config_file,
+            device=device,
+            auto_collate=auto_collate,
+            compile=kwargs.pop('compile', False),
+            compile_options=kwargs.pop('compile_options', {}))
+
+        assert isinstance(self.model, Model), \
+            f'please check whether model config exists in {ModelFile.CONFIGURATION}'
+
+        if preprocessor is None:
+            self.preprocessor = Preprocessor.from_pretrained(
+                self.model.model_dir,
+                sequence_length=sequence_length,
+                **kwargs)
+        self.model.eval()
+
+        assert hasattr(self.preprocessor, 'id2label')
+        self.id2label = self.preprocessor.id2label
+
+    def forward(self, inputs: Dict[str, Any],
+                **forward_params) -> Dict[str, Any]:
+        text = inputs.pop(OutputKeys.TEXT)
+        with torch.no_grad():
+            outputs = self.model(**inputs, **forward_params)
+            return {**outputs, OutputKeys.TEXT: text}
+
+    def postprocess(self, inputs: Dict[str, Any],
+                    **postprocess_params) -> Dict[str, Any]:
+        r"""Precess the prediction results
+            Args:
+                inputs (dict[str, Any]): should be tensors from model
+
+            Returns:
+                Dict[str, Any]: the prediction results
+        """
+        text = inputs['text']
+
+        if OutputKeys.PREDICTIONS not in inputs:
+            logits = inputs[OutputKeys.LOGITS]
+            if len(logits.shape) == 3:
+                logits = logits[0]
+            predictions = torch.argmax(logits, dim=-1)
+        else:
+            predictions = inputs[OutputKeys.PREDICTIONS]
+            if len(predictions.shape) == 2:
+                predictions = predictions[0]
+
+        binary_prediction = []
+        for i, ch in enumerate(text):
+            if ch in self.PUNC_LIST:
+                binary_prediction.append(0)
+            else:
+                binary_prediction.append(-100)
+
+        result_text = ''
+        for i, p in enumerate(predictions):
+            if i >= len(text):
+                continue
+            result_text += text[i]
+            if binary_prediction[i] != -100:
+                binary_prediction[i] = p
+                if p == 1:
+                    result_text += '|'
+
+        outputs = {
+            'text': inputs['text'],
+            'logits': inputs['logits'],
+            'prediction': binary_prediction
+        }
+        return outputs

modelscope/preprocessors/base.py

@@ -48,6 +48,8 @@ PREPROCESSOR_MAP = {
     Preprocessors.token_cls_tokenizer,
     (Models.bert, Tasks.token_classification):
     Preprocessors.token_cls_tokenizer,
+    (Models.bert, Tasks.speaker_diarization_semantic_speaker_turn_detection):
+    Preprocessors.token_cls_tokenizer,
     (Models.bert, Tasks.word_segmentation):
     Preprocessors.token_cls_tokenizer,
 

modelscope/preprocessors/speaker.py

@@ -2,18 +2,20 @@
 
 from typing import Any, Dict, List, Tuple, Union
 
+import torch
+
 from modelscope.metainfo import Preprocessors
 from modelscope.preprocessors import Preprocessor
 from modelscope.preprocessors.builder import PREPROCESSORS
 from modelscope.preprocessors.nlp.text_classification_preprocessor import \
     TextClassificationPreprocessorBase
+from modelscope.preprocessors.nlp.token_classification_preprocessor import (
+    NLPTokenizerForLSTM, TokenClassificationPreprocessorBase)
 from modelscope.preprocessors.nlp.transformers_tokenizer import NLPTokenizer
 from modelscope.utils.constant import Fields, ModeKeys
 from modelscope.utils.hub import get_model_type, parse_label_mapping
 from modelscope.utils.logger import get_logger
 
-# from modelscope.preprocessors.nlp.utils import labels_to_id, parse_text_and_label
-
 logger = get_logger()
 
 
@@ -52,3 +54,261 @@ class SpeakerDiarizationDialogueDetectionPreprocessor(
             model_dir, model_type, use_fast=use_fast, tokenize_kwargs=kwargs)
         super().__init__(model_dir, first_sequence, second_sequence, label,
                          label2id, mode, keep_original_columns)
+
+
+@PREPROCESSORS.register_module(
+    Fields.audio, module_name=Preprocessors.token_cls_tokenizer)
+class SpeakerDiarizationSemanticSpeakerTurnDetectionPreprocessor(
+        TokenClassificationPreprocessorBase):
+
+    def __init__(self,
+                 model_dir: str = None,
+                 first_sequence: str = 'text',
+                 label: str = 'label',
+                 label2id: Dict = None,
+                 label_all_tokens: bool = False,
+                 mode: str = ModeKeys.INFERENCE,
+                 max_length=None,
+                 use_fast=None,
+                 keep_original_columns=None,
+                 return_text=True,
+                 **kwargs):
+        super().__init__(model_dir, first_sequence, label, label2id,
+                         label_all_tokens, mode, keep_original_columns,
+                         return_text)
+        model_type = None
+        if model_dir is not None:
+            model_type = get_model_type(model_dir)
+
+        kwargs['truncation'] = kwargs.get('truncation', True)
+        kwargs['padding'] = kwargs.get('padding', 'max_length')
+        kwargs[
+            'max_length'] = max_length if max_length is not None else kwargs.get(
+                'sequence_length', 128)
+        kwargs.pop('sequence_length', None)
+        kwargs['add_special_tokens'] = model_type != 'lstm'
+        self.nlp_tokenizer = NLPTokenizerForLSTM(
+            model_dir=model_dir,
+            model_type=model_type,
+            use_fast=use_fast,
+            tokenize_kwargs=kwargs)
+
+    def _tokenize_text(self, text: Union[str, List[str]], **kwargs):
+        tokens = text
+
+        if self.mode != ModeKeys.INFERENCE:
+            assert isinstance(tokens, list), 'Input needs to be lists in training and evaluating,' \
+                'because the length of the words and the labels need to be equal.'
+        is_split_into_words = self.nlp_tokenizer.get_tokenizer_kwarg(
+            'is_split_into_words', False)
+        if is_split_into_words:
+            # for supporting prompt seperator, should split twice. [SEP] for default.
+            sep_idx = tokens.find('[SEP]')
+            if sep_idx == -1 or self.is_lstm_model:
+                tokens = list(tokens)
+            else:
+                tmp_tokens = []
+                tmp_tokens.extend(list(tokens[:sep_idx]))
+                tmp_tokens.append('[SEP]')
+                tmp_tokens.extend(list(tokens[sep_idx + 5:]))
+                tokens = tmp_tokens
+
+        if is_split_into_words and self.mode == ModeKeys.INFERENCE:
+            encodings, word_ids = self._tokenize_text_by_words(
+                tokens, **kwargs)
+        elif self.nlp_tokenizer.tokenizer.is_fast:
+            encodings, word_ids = self._tokenize_text_with_fast_tokenizer(
+                tokens, **kwargs)
+        else:
+            encodings, word_ids = self._tokenize_text_with_slow_tokenizer(
+                tokens, **kwargs)
+
+        sep_idx = -1
+        for idx, token_id in enumerate(encodings['input_ids']):
+            if token_id == self.nlp_tokenizer.tokenizer.sep_token_id:
+                sep_idx = idx
+                break
+        if sep_idx != -1:
+            for i in range(sep_idx, len(encodings['label_mask'])):
+                encodings['label_mask'][i] = False
+
+        if self.mode == ModeKeys.INFERENCE:
+            for key in encodings.keys():
+                encodings[key] = torch.tensor(encodings[key]).unsqueeze(0)
+        else:
+            encodings.pop('offset_mapping', None)
+        return encodings, word_ids
+
+    def _tokenize_text_by_words(self, tokens, **kwargs):
+        input_ids = []
+        label_mask = []
+        offset_mapping = []
+        attention_mask = []
+
+        for offset, token in enumerate(tokens):
+            subtoken_ids = self.nlp_tokenizer.tokenizer.encode(
+                token, add_special_tokens=False)
+            if len(subtoken_ids) == 0:
+                subtoken_ids = [self.nlp_tokenizer.tokenizer.unk_token_id]
+            input_ids.extend(subtoken_ids)
+            attention_mask.extend([1] * len(subtoken_ids))
+            label_mask.extend([True] + [False] * (len(subtoken_ids) - 1))
+            offset_mapping.extend([(offset, offset + 1)])
+
+        padding = kwargs.get('padding',
+                             self.nlp_tokenizer.get_tokenizer_kwarg('padding'))
+        max_length = kwargs.get(
+            'max_length',
+            kwargs.get('sequence_length',
+                       self.nlp_tokenizer.get_tokenizer_kwarg('max_length')))
+        special_token = 1 if self.nlp_tokenizer.get_tokenizer_kwarg(
+            'add_special_tokens') else 0
+        if len(label_mask) > max_length - 2 * special_token:
+            label_mask = label_mask[:(max_length - 2 * special_token)]
+            input_ids = input_ids[:(max_length - 2 * special_token)]
+        offset_mapping = offset_mapping[:sum(label_mask)]
+        if padding == 'max_length':
+            label_mask = [False] * special_token + label_mask + \
+                [False] * (max_length - len(label_mask) - special_token)
+            offset_mapping = offset_mapping + [(0, 0)] * (
+                max_length - len(offset_mapping))
+            input_ids = [self.nlp_tokenizer.tokenizer.cls_token_id] * special_token + input_ids + \
+                [self.nlp_tokenizer.tokenizer.sep_token_id] * special_token + \
+                [self.nlp_tokenizer.tokenizer.pad_token_id] * (max_length - len(input_ids) - 2 * special_token)
+            attention_mask = attention_mask + [1] * (
+                special_token * 2) + [0] * (
+                    max_length - len(attention_mask) - 2 * special_token)
+        else:
+            label_mask = [False] * special_token + label_mask + \
+                [False] * special_token
+            input_ids = [self.nlp_tokenizer.tokenizer.cls_token_id] * special_token + input_ids + \
+                [self.nlp_tokenizer.tokenizer.sep_token_id] * special_token
+            attention_mask = attention_mask + [1] * (special_token * 2)
+
+        encodings = {
+            'input_ids': input_ids,
+            'attention_mask': attention_mask,
+            'label_mask': label_mask,
+            'offset_mapping': offset_mapping,
+        }
+        return encodings, None
+
+    def _tokenize_text_with_fast_tokenizer(self, tokens, **kwargs):
+        is_split_into_words = isinstance(tokens, list)
+        encodings = self.nlp_tokenizer(
+            tokens,
+            return_offsets_mapping=True,
+            is_split_into_words=is_split_into_words,
+            **kwargs)
+        label_mask = []
+        word_ids = encodings.word_ids()
+        offset_mapping = []
+        for i in range(len(word_ids)):
+            if word_ids[i] is None:
+                label_mask.append(False)
+            elif word_ids[i] == word_ids[i - 1]:
+                label_mask.append(False)
+                if not is_split_into_words:
+                    offset_mapping[-1] = (offset_mapping[-1][0],
+                                          encodings['offset_mapping'][i][1])
+            else:
+                label_mask.append(True)
+                if is_split_into_words:
+                    offset_mapping.append((word_ids[i], word_ids[i] + 1))
+                else:
+                    offset_mapping.append(encodings['offset_mapping'][i])
+
+        padding = self.nlp_tokenizer.get_tokenizer_kwarg('padding')
+        if padding == 'max_length':
+            offset_mapping = offset_mapping + [(0, 0)] * (
+                len(label_mask) - len(offset_mapping))
+        encodings['offset_mapping'] = offset_mapping
+        encodings['label_mask'] = label_mask
+        return encodings, word_ids
+
+    def _tokenize_text_with_slow_tokenizer(self, tokens, **kwargs):
+        assert self.mode == ModeKeys.INFERENCE and isinstance(tokens, str), \
+            'Slow tokenizer now only support str input in inference mode. If you are training models, ' \
+            'please consider using the fast tokenizer.'
+        word_ids = None
+        encodings = self.nlp_tokenizer(
+            tokens, is_split_into_words=False, **kwargs)
+        tokenizer_name = self.nlp_tokenizer.get_tokenizer_class()
+        method = 'get_label_mask_and_offset_mapping_' + tokenizer_name
+        if not hasattr(self, method):
+            raise RuntimeError(
+                f'No `{method}` method defined for '
+                f'tokenizer {tokenizer_name}, please use a fast tokenizer instead, or '
+                f'try to implement a `{method}` method')
+        label_mask, offset_mapping = getattr(self, method)(tokens)
+        padding = kwargs.get('padding',
+                             self.nlp_tokenizer.get_tokenizer_kwarg('padding'))
+        max_length = kwargs.get(
+            'max_length', self.nlp_tokenizer.get_tokenizer_kwarg('max_length'))
+        special_token = 1 if kwargs.get(
+            'add_special_tokens',
+            self.nlp_tokenizer.get_tokenizer_kwarg(
+                'add_special_tokens')) else 0
+        if len(label_mask) > max_length - 2 * special_token:
+            label_mask = label_mask[:(max_length - 2 * special_token)]
+        offset_mapping = offset_mapping[:sum(label_mask)]
+
+        if padding == 'max_length':
+            label_mask = [False] * special_token + label_mask + \
+                [False] * (max_length - len(label_mask) - special_token)
+            offset_mapping = offset_mapping + [(0, 0)] * (
+                max_length - len(offset_mapping))
+        else:
+            label_mask = [False] * special_token + label_mask + \
+                [False] * special_token
+        encodings['offset_mapping'] = offset_mapping
+        encodings['label_mask'] = label_mask
+        return encodings, word_ids
+
+    def get_label_mask_and_offset_mapping_BertTokenizer(self, text):
+        label_mask = []
+        offset_mapping = []
+        tokens = self.nlp_tokenizer.tokenizer.tokenize(text)
+        offset = 0
+        for token in tokens:
+            is_start = (token[:2] != '##')
+            if is_start:
+                label_mask.append(True)
+            else:
+                token = token[2:]
+                label_mask.append(False)
+            start = offset + text[offset:].index(token)
+            end = start + len(token)
+            if is_start:
+                offset_mapping.append((start, end))
+            else:
+                offset_mapping[-1] = (offset_mapping[-1][0], end)
+            offset = end
+
+        return label_mask, offset_mapping
+
+    def get_label_mask_and_offset_mapping_XLMRobertaTokenizer(self, text):
+        label_mask = []
+        offset_mapping = []
+        tokens = self.nlp_tokenizer.tokenizer.tokenize(text)
+        offset = 0
+        last_is_blank = False
+        for token in tokens:
+            is_start = (token[0] == '_')
+            if is_start:
+                token = token[1:]
+                label_mask.append(True)
+                if len(token) == 0:
+                    last_is_blank = True
+                    continue
+            else:
+                label_mask.append(False)
+            start = offset + text[offset:].index(token)
+            end = start + len(token)
+            if last_is_blank or is_start:
+                offset_mapping.append((start, end))
+            else:
+                offset_mapping[-1] = (offset_mapping[-1][0], end)
+            offset = end
+            last_is_blank = False
+        return label_mask, offset_mapping

modelscope/utils/constant.py

@@ -229,6 +229,7 @@ class AudioTasks(object):
     language_score_prediction = 'language-score-prediction'
     speech_timestamp = 'speech-timestamp'
     speaker_diarization_dialogue_detection = 'speaker-diarization-dialogue-detection'
+    speaker_diarization_semantic_speaker_turn_detection = 'speaker-diarization-semantic-speaker-turn-detection'
 
 
 class MultiModalTasks(object):

tests/pipelines/test_speaker_diarization_semantic_speaker_turn_detection.py

@@ -0,0 +1,62 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import unittest
+from typing import Any, Dict, Optional
+
+from modelscope.pipelines import pipeline
+from modelscope.utils.constant import Tasks
+from modelscope.utils.logger import get_logger
+from modelscope.utils.test_utils import test_level
+
+logger = get_logger()
+
+
+class SpeakerDiarizationSemanticSpeakerTurnDetectionTest(unittest.TestCase):
+
+    test_datasets = [{
+        'sentence': '嗯，到时候有问题我再跟您联系吧，刘老师。行，可以的，那到时候再联系吧。',
+    }, {
+        'sentence':
+        '你是如何看待这个问题的呢？这个问题挺好解决的，我们只需要增加停车位就行了。嗯嗯，好，那我们业主就放心了。'
+    }, {
+        'sentence': '这个电台播放各种音乐，包括古典、爵士、民族等各种风格，并附有专业的音乐解说。'
+    }]
+
+    semantic_std_model_id = 'damo/speech_bert_semantic-spk-turn-detection-punc_speaker-diarization_chinese'
+
+    def setUp(self) -> None:
+        self.task = Tasks.speaker_diarization_semantic_speaker_turn_detection
+
+    def run_pipeline(self,
+                     model_id: str,
+                     model_revision=None) -> Dict[str, Any]:
+        speaker_turn_detection = pipeline(
+            task=self.task, model=model_id, model_revision=model_revision)
+        output_list = []
+        for sentence_item in self.test_datasets:
+            sentence = sentence_item['sentence']
+            output_list.append((sentence, speaker_turn_detection(sentence)))
+        return output_list
+
+    @unittest.skipUnless(test_level() >= 0, 'Skip test in current test level')
+    def test_semantic_speaker_turn_detection_model(self):
+        logger.info('Run speaker diarization speaker turn detection')
+
+        pipeline_results = self.run_pipeline(
+            model_id=self.semantic_std_model_id, model_revision='v0.5.0')
+        for sentence, result in pipeline_results:
+            cur_predict_sentence = ''
+            predict = result['prediction']
+            for i, ch in enumerate(sentence):
+                cur_predict_sentence += ch
+                if i >= len(predict):
+                    continue
+                if predict[i] == 1:
+                    cur_predict_sentence += '|'
+            logger.info(f'\nresult = {result.keys()}'
+                        f'\nsentence = {sentence}'
+                        f'\npredict  = {cur_predict_sentence}')
+        logger.info('Text semantic_speaker_turn_detection model finished')
+
+
+if __name__ == '__main__':
+    unittest.main()