diff --git a/modelscope/models/nlp/task_models/token_classification.py b/modelscope/models/nlp/task_models/token_classification.py index aa84eaf0..8c5142b9 100644 --- a/modelscope/models/nlp/task_models/token_classification.py +++ b/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, diff --git a/modelscope/pipelines/nlp/token_classification_pipeline.py b/modelscope/pipelines/nlp/token_classification_pipeline.py index 9fd8e325..0c87e3a0 100644 --- a/modelscope/pipelines/nlp/token_classification_pipeline.py +++ b/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 diff --git a/tests/pipelines/plugin_remote_pipelines/test_plugin_model.py b/tests/pipelines/plugin_remote_pipelines/test_plugin_model.py index 71b9e64f..aeb6c9bd 100644 --- a/tests/pipelines/plugin_remote_pipelines/test_plugin_model.py +++ b/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' + ) diff --git a/tests/pipelines/test_named_entity_recognition.py b/tests/pipelines/test_named_entity_recognition.py index 8b7424f4..4f431b9f 100644 --- a/tests/pipelines/test_named_entity_recognition.py +++ b/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: