AutoAndroidController/py/venv/Lib/site-packages/modelscope/pipelines/nlp/document_segmentation_pipeline.py

# Copyright (c) Alibaba, Inc. and its affiliates.

import re
from typing import Any, Dict, List, Union

import numpy as np
import torch
from datasets import Dataset

from modelscope.metainfo import Pipelines
from modelscope.models import Model
from modelscope.outputs import OutputKeys
from modelscope.pipelines.base import Pipeline, Tensor
from modelscope.pipelines.builder import PIPELINES
from modelscope.preprocessors import \
    DocumentSegmentationTransformersPreprocessor
from modelscope.utils.constant import Tasks
from modelscope.utils.logger import get_logger

logger = get_logger()

__all__ = ['DocumentSegmentationPipeline']


@PIPELINES.register_module(
    Tasks.document_segmentation, module_name=Pipelines.document_segmentation)
class DocumentSegmentationPipeline(Pipeline):

    def __init__(
            self,
            model: Union[Model, str],
            preprocessor: DocumentSegmentationTransformersPreprocessor = None,
            config_file: str = None,
            device: str = 'gpu',
            auto_collate=True,
            **kwargs):
        """The document segmentation pipeline.

        Args:
            model (str or Model): Supply either a local model dir or a model id from the model hub
            preprocessor (Preprocessor): An optional preprocessor instance, please make sure the preprocessor fits for
            the model if supplied.
        """
        super().__init__(
            model=model,
            preprocessor=preprocessor,
            config_file=config_file,
            device=device,
            auto_collate=auto_collate,
            **kwargs)

        kwargs.pop('compile', None)
        kwargs.pop('compile_options', None)

        self.model_dir = self.model.model_dir
        self.model_cfg = self.model.model_cfg
        if preprocessor is None:
            self.preprocessor = DocumentSegmentationTransformersPreprocessor(
                self.model_dir, self.model.config.max_position_embeddings,
                **kwargs)

    def __call__(
            self, documents: Union[List[List[str]], List[str],
                                   str]) -> Dict[str, Any]:
        output = self.predict(documents)
        output = self.postprocess(output)
        return output

    def predict(
            self, documents: Union[List[List[str]], List[str],
                                   str]) -> Dict[str, Any]:
        pred_samples = self.cut_documents(documents)

        if self.model_cfg['level'] == 'topic':
            paragraphs = pred_samples.pop('paragraphs')

        predict_examples = Dataset.from_dict(pred_samples)

        # Predict Feature Creation
        predict_dataset = self.preprocessor(predict_examples, self.model_cfg)
        num_examples = len(
            predict_examples[self.preprocessor.context_column_name])
        num_samples = len(
            predict_dataset[self.preprocessor.context_column_name])

        if self.model_cfg['type'] == 'bert':
            predict_dataset.pop('segment_ids')

        labels = predict_dataset.pop('labels')
        sentences = predict_dataset.pop('sentences')
        example_ids = predict_dataset.pop(
            self.preprocessor.example_id_column_name)

        if (self.model or (self.has_multiple_models and self.models[0])):
            if not self._model_prepare:
                self.prepare_model()

        with torch.no_grad():
            input = {
                key: torch.tensor(val).to(self.device)
                for key, val in predict_dataset.items()
            }
            predictions = self.model.forward(**input).logits.cpu()

        predictions = np.argmax(predictions, axis=2)
        assert len(sentences) == len(
            predictions), 'sample {}  infer_sample {} prediction {}'.format(
                num_samples, len(sentences), len(predictions))
        # Remove ignored index (special tokens)

        true_predictions = [
            [
                self.preprocessor.label_list[p]
                for (p, l) in zip(prediction, label) if l != -100  # noqa *
            ] for prediction, label in zip(predictions, labels)
        ]

        true_labels = [
            [
                self.preprocessor.label_list[l]
                for (p, l) in zip(prediction, label) if l != -100  # noqa *
            ] for prediction, label in zip(predictions, labels)
        ]

        # Save predictions
        out = []
        for i in range(num_examples):
            if self.model_cfg['level'] == 'topic':
                out.append({
                    'sentences': [],
                    'labels': [],
                    'predictions': [],
                    'paragraphs': paragraphs[i]
                })
            else:
                out.append({'sentences': [], 'labels': [], 'predictions': []})

        for prediction, sentence_list, label, example_id in zip(
                true_predictions, sentences, true_labels, example_ids):
            if self.model_cfg['level'] == 'doc':
                if len(label) < len(sentence_list):
                    label.append('B-EOP')
                    prediction.append('B-EOP')
                assert len(sentence_list) == len(prediction), '{} {}'.format(
                    len(sentence_list), len(prediction))
                assert len(sentence_list) == len(label), '{} {}'.format(
                    len(sentence_list), len(label))

            out[example_id]['sentences'].extend(sentence_list)
            out[example_id]['labels'].extend(label)
            out[example_id]['predictions'].extend(prediction)

        if self.model_cfg['level'] == 'topic':
            for i in range(num_examples):
                assert len(out[i]['predictions']) + 1 == len(
                    out[i]['paragraphs'])
                out[i]['predictions'].append('B-EOP')
                out[i]['labels'].append('B-EOP')

        return out

    def postprocess(self, inputs: Dict[str, Tensor]) -> Dict[str, Tensor]:
        """process the prediction results

        Args:
            inputs (Dict[str, Any]): _description_

        Returns:
            Dict[str, str]: the prediction results
        """
        result = []
        res_preds = []
        list_count = len(inputs)

        if self.model_cfg['level'] == 'topic':
            for num in range(list_count):
                res = []
                pred = []
                for s, p, l in zip(inputs[num]['paragraphs'],
                                   inputs[num]['predictions'],
                                   inputs[num]['labels']):
                    s = s.strip()
                    if p == 'B-EOP':
                        s = ''.join([s, '\n\n\t'])
                        pred.append(1)
                    else:
                        s = ''.join([s, '\n\t'])
                        pred.append(0)
                    res.append(s)
                res_preds.append(pred)
                document = ('\t' + ''.join(res).strip())
                result.append(document)
        else:
            for num in range(list_count):
                res = []
                for s, p in zip(inputs[num]['sentences'],
                                inputs[num]['predictions']):
                    s = s.strip()
                    if p == 'B-EOP':
                        s = ''.join([s, '\n\t'])
                    res.append(s)

                document = ('\t' + ''.join(res))
                result.append(document)

        if list_count == 1:
            return {OutputKeys.TEXT: result[0]}
        else:
            return {OutputKeys.TEXT: result}

    def cut_documents(self, para: Union[List[List[str]], List[str], str]):
        document_list = para
        paragraphs = []
        sentences = []
        labels = []
        example_id = []
        id = 0

        if self.model_cfg['level'] == 'topic':
            if isinstance(para, str):
                document_list = [[para]]
            elif isinstance(para[0], str):
                document_list = [para]

            for document in document_list:
                sentence = []
                label = []
                for item in document:
                    sentence_of_current_paragraph = self.cut_sentence(item)
                    sentence.extend(sentence_of_current_paragraph)
                    label.extend(['-100']
                                 * (len(sentence_of_current_paragraph) - 1)
                                 + ['B-EOP'])
                paragraphs.append(document)
                sentences.append(sentence)
                labels.append(label)
                example_id.append(id)
                id += 1

            return {
                'example_id': example_id,
                'sentences': sentences,
                'paragraphs': paragraphs,
                'labels': labels
            }
        else:
            if isinstance(para, str):
                document_list = [para]

            for document in document_list:
                sentence = self.cut_sentence(document)
                label = ['O'] * (len(sentence) - 1) + ['B-EOP']
                sentences.append(sentence)
                labels.append(label)
                example_id.append(id)
                id += 1

            return {
                'example_id': example_id,
                'sentences': sentences,
                'labels': labels
            }

    def cut_sentence(self, para):
        para = re.sub(r'([。！.!？\?])([^”’])', r'\1\n\2', para)  # noqa *
        para = re.sub(r'(\.{6})([^”’])', r'\1\n\2', para)  # noqa *
        para = re.sub(r'(\…{2})([^”’])', r'\1\n\2', para)  # noqa *
        para = re.sub(r'([。！？\?][”’])([^，。！？\?])', r'\1\n\2', para)  # noqa *
        para = para.rstrip()
        return [_ for _ in para.split('\n') if _]