enable token_cls_pipeline to inference on longer inputs and return entity probabilities (#551)

* allow token classification pipelines to predict longer sentences

* bugfix

* skip adaseq pipeline ut when connection error occurs

* return entity probabilities

modelscope/models/nlp/task_models/token_classification.py

@@ -102,6 +102,7 @@ class ModelForTokenClassificationWithCRF(ModelForTokenClassification):
     base_model_prefix = 'encoder'
 
     def postprocess(self, inputs, **kwargs):
+        logits = inputs['logits']
         predicts = self.head.decode(inputs['logits'], inputs['label_mask'])
         offset_mapping = inputs['offset_mapping']
         mask = inputs['label_mask']
@@ -119,7 +120,7 @@ class ModelForTokenClassificationWithCRF(ModelForTokenClassification):
 
         return AttentionTokenClassificationModelOutput(
             loss=None,
-            logits=None,
+            logits=logits,
             hidden_states=None,
             attentions=None,
             label_mask=mask,

modelscope/pipelines/nlp/token_classification_pipeline.py

@@ -1,6 +1,7 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 
-from typing import Any, Dict, List, Optional, Union
+import math
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -8,7 +9,7 @@ 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.base import Input, Pipeline
 from modelscope.pipelines.builder import PIPELINES
 from modelscope.preprocessors import Preprocessor
 from modelscope.utils.constant import ModelFile, Tasks
@@ -64,6 +65,7 @@ class TokenClassificationPipeline(Pipeline):
                 sequence_length=sequence_length,
                 **kwargs)
         self.model.eval()
+        self.sequence_length = sequence_length
 
         assert hasattr(self.preprocessor, 'id2label')
         self.id2label = self.preprocessor.id2label
@@ -131,9 +133,20 @@ class TokenClassificationPipeline(Pipeline):
         predictions = torch_nested_numpify(torch_nested_detach(predictions))
         labels = [self.id2label[x] for x in predictions]
 
+        return_prob = postprocess_params.pop('return_prob', True)
+        if return_prob:
+            if OutputKeys.LOGITS in inputs:
+                logits = inputs[OutputKeys.LOGITS]
+                if len(logits.shape) == 3:
+                    logits = logits[0]
+                probs = torch_nested_numpify(
+                    torch_nested_detach(logits.softmax(-1)))
+            else:
+                return_prob = False
+
         chunks = []
         chunk = {}
-        for label, offsets in zip(labels, offset_mapping):
+        for i, (label, offsets) in enumerate(zip(labels, offset_mapping)):
             if label[0] in 'BS':
                 if chunk:
                     chunk['span'] = text[chunk['start']:chunk['end']]
@@ -143,6 +156,8 @@ class TokenClassificationPipeline(Pipeline):
                     'start': offsets[0],
                     'end': offsets[1]
                 }
+                if return_prob:
+                    chunk['prob'] = probs[i][predictions[i]]
             if label[0] in 'I':
                 if not chunk:
                     chunk = {
@@ -150,6 +165,8 @@ class TokenClassificationPipeline(Pipeline):
                         'start': offsets[0],
                         'end': offsets[1]
                     }
+                    if return_prob:
+                        chunk['prob'] = probs[i][predictions[i]]
             if label[0] in 'E':
                 if not chunk:
                     chunk = {
@@ -157,6 +174,8 @@ class TokenClassificationPipeline(Pipeline):
                         'start': offsets[0],
                         'end': offsets[1]
                     }
+                    if return_prob:
+                        chunk['prob'] = probs[i][predictions[i]]
             if label[0] in 'IES':
                 if chunk:
                     chunk['end'] = offsets[1]
@@ -172,3 +191,63 @@ class TokenClassificationPipeline(Pipeline):
             chunks.append(chunk)
 
         return chunks
+
+    def _process_single(self, input: Input, *args, **kwargs) -> Dict[str, Any]:
+        split_max_length = kwargs.pop('split_max_length',
+                                      0)  # default: no split
+        if split_max_length <= 0:
+            return super()._process_single(input, *args, **kwargs)
+        else:
+            split_texts, index_mapping = self._auto_split([input],
+                                                          split_max_length)
+            outputs = []
+            for text in split_texts:
+                outputs.append(super()._process_single(text, *args, **kwargs))
+            return self._auto_join(outputs, index_mapping)[0]
+
+    def _process_batch(self, input: List[Input], batch_size: int, *args,
+                       **kwargs) -> List[Dict[str, Any]]:
+        split_max_length = kwargs.pop('split_max_length',
+                                      0)  # default: no split
+        if split_max_length <= 0:
+            return super()._process_batch(
+                input, batch_size=batch_size, *args, **kwargs)
+        else:
+            split_texts, index_mapping = self._auto_split(
+                input, split_max_length)
+            outputs = super()._process_batch(
+                split_texts, batch_size=batch_size, *args, **kwargs)
+            return self._auto_join(outputs, index_mapping)
+
+    def _auto_split(self, input_texts: List[str], split_max_length: int):
+        split_texts = []
+        index_mapping = {}
+        new_idx = 0
+        for raw_idx, text in enumerate(input_texts):
+            if len(text) < split_max_length:
+                split_texts.append(text)
+                index_mapping[new_idx] = (raw_idx, 0)
+                new_idx += 1
+            else:
+                n_split = math.ceil(len(text) / split_max_length)
+                for i in range(n_split):
+                    offset = i * split_max_length
+                    split_texts.append(text[offset:offset + split_max_length])
+                    index_mapping[new_idx] = (raw_idx, offset)
+                    new_idx += 1
+        return split_texts, index_mapping
+
+    def _auto_join(
+            self, outputs: List[Dict[str, Any]],
+            index_mapping: Dict[int, Tuple[int, int]]) -> List[Dict[str, Any]]:
+        joined_outputs = []
+        for idx, output in enumerate(outputs):
+            raw_idx, offset = index_mapping[idx]
+            if raw_idx >= len(joined_outputs):
+                joined_outputs.append(output)
+            else:
+                for chunk in output[OutputKeys.OUTPUT]:
+                    chunk['start'] += offset
+                    chunk['end'] += offset
+                    joined_outputs[raw_idx][OutputKeys.OUTPUT].append(chunk)
+        return joined_outputs

tests/pipelines/plugin_remote_pipelines/test_plugin_model.py

@@ -23,20 +23,31 @@ class PluginModelTest(unittest.TestCase):
 
     @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
     def test_run_span_based_ner_pipeline(self):
-        pipeline_ins = pipeline(
-            Tasks.named_entity_recognition,
-            'damo/nlp_nested-ner_named-entity-recognition_chinese-base-med')
-        print(
-            pipeline_ins(
-                '1、可测量目标： 1周内胸闷缓解。2、下一步诊疗措施：1.心内科护理常规，一级护理，低盐低脂饮食，留陪客。'
-                '2.予“阿司匹林肠溶片”抗血小板聚集，“呋塞米、螺内酯”利尿减轻心前负荷，“瑞舒伐他汀”调脂稳定斑块，“厄贝沙坦片片”降血压抗心机重构'
-            ))
+        try:
+            pipeline_ins = pipeline(
+                Tasks.named_entity_recognition,
+                'damo/nlp_nested-ner_named-entity-recognition_chinese-base-med'
+            )
+            print(
+                pipeline_ins(
+                    '1、可测量目标： 1周内胸闷缓解。2、下一步诊疗措施：1.心内科护理常规，一级护理，低盐低脂饮食，留陪客。'
+                    '2.予“阿司匹林肠溶片”抗血小板聚集，“呋塞米、螺内酯”利尿减轻心前负荷，“瑞舒伐他汀”调脂稳定斑块，“厄贝沙坦片片”降血压抗心机重构'
+                ))
+        except RuntimeError:
+            print(
+                'Skip test span_based_ner_pipeline! RuntimeError: Try loading from huggingface and modelscope failed'
+            )
 
     def test_maoe_pipelines(self):
-        pipeline_ins = pipeline(
-            Tasks.named_entity_recognition,
-            'damo/nlp_maoe_named-entity-recognition_chinese-base-general')
-        print(
-            pipeline_ins(
-                '刘培强，男，生理年龄40岁（因为在太空中进入休眠状态），实际年龄52岁，领航员国际空间站中的中国航天员，机械工程专家，军人，军衔中校。'
-            ))
+        try:
+            pipeline_ins = pipeline(
+                Tasks.named_entity_recognition,
+                'damo/nlp_maoe_named-entity-recognition_chinese-base-general')
+            print(
+                pipeline_ins(
+                    '刘培强，男，生理年龄40岁（因为在太空中进入休眠状态），实际年龄52岁，领航员国际空间站中的中国航天员，机械工程专家，军人，军衔中校。'
+                ))
+        except RuntimeError:
+            print(
+                'Skip test maoe_pipeline! RuntimeError: Try loading from huggingface and modelscope failed'
+            )

tests/pipelines/test_named_entity_recognition.py

@@ -459,6 +459,25 @@ class NamedEntityRecognitionTest(unittest.TestCase):
         pipeline_ins = pipeline(task=Tasks.named_entity_recognition)
         print(pipeline_ins(input=self.sentence))
 
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_run_long_chinese_with_model_name(self):
+        pipeline_ins = pipeline(
+            task=Tasks.named_entity_recognition, model=self.chinese_model_id)
+        print(
+            pipeline_ins(
+                input=self.sentence + '. ' * 1000,
+                split_max_length=300))  # longer than 512
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_run_long_chinese_with_model_name_batch(self):
+        pipeline_ins = pipeline(
+            task=Tasks.named_entity_recognition, model=self.chinese_model_id)
+        print(
+            pipeline_ins(
+                input=[self.sentence + '. ' * 1000] * 2,
+                batch_size=2,
+                split_max_length=300))  # longer than 512
+
     @unittest.skipUnless(test_level() >= 2, 'skip test in current test level')
     def test_run_with_all_modelcards(self):
         for item in self.all_modelcards_info: