modelscope/tests/pipelines/test_extractive_summarization.py

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

import unittest
from typing import Any, Dict

from modelscope.outputs import OutputKeys
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 ExtractiveSummarizationTest(unittest.TestCase):

    def setUp(self) -> None:
        self.task = Tasks.extractive_summarization

    ponet_doc_model_id = 'damo/nlp_ponet_extractive-summarization_doc-level_chinese-base'
    ponet_topic_model_id = 'damo/nlp_ponet_extractive-summarization_topic-level_chinese-base'

    sentences = '嗨吃晚饭了吗？大家都，哎。嗯。还没吃呢？吃了。哎不好意思，这么晚把大家叫过来，商量个事儿啊，下个月咱们学校要开个文艺晚会，哎挺挺赶，时间也挺赶的，想把大家叫过来商量一下。没事儿。啊。对，校长校长这人太不靠谱儿。嗯。咋了他。咋。每次都是，这种大的活动之前他就出去，他就出出出差什么之类。他可能想逃避讲话什么的。不懂每个学校都是这样儿的吗？对。嗯可能有可能。每个学校的校长都是这样儿的吗？每次讲的话还都是一样的他那。对，说不定就这样儿呢。十分十分枯燥。唉领导就是忙，那怎么办？联系一下呗，看他是有时间，来一趟或者是录个小视频什么的，是不是。肯定要让他回来呀，他不回来这个文艺晚会举办不了的，得要他发言。嗯。那就提前打电话联系一下。我可联系不上，不知道你们谁能联系上。我觉得这个就交给你了，比较好。行吧，那我哪天去他办公室，看看他吧。但是他肯定是要讲话什么的呀？嗯是他肯定是要讲话，那除了他，还有谁能讲话呀？一些。嗯，那种院的什么辅导员啊或者什么。对。其实不止校长，院长也可以安排一下。或者一些主一些主任什么都可以。院长也可以，院长也可以，主任也可以，对。嗯。哎，他们是在台上坐一排的那种吗？台上坐一排，那有点过分了吧。别别别那样大家都放不开啦。应该坐在台下吧。这毕竟大家还要表演嘛，是不是。嗨不起来了。就舞台底下可以坐一排，然后摆一些牌啊，然后把他们叫什么都摆一下。对，他。那我去。对让他们坐在第一排挺好。我觉得文艺晚会主要是学生看，那领导那一排那挡住了怎么办。就充个场子样子。嗯。他们可能也不看啊，就刚开始然后来然后演节目他们好像就走了。就是，嗯就是看两眼吧，出个面。那既然。既然这样儿，咱们学生加上领导那是不是要再请些家长之类的。不但我们觉得还要商量一下，就是这个领导座位的顺序，就是谁该坐在中间，还坐在边上什么的。那这个。我觉得领导可以最后出场。那。这个肯定校校长是最，要坐在中间的呀。但刚开始不都是就是，欢迎领导然后入场啊什么，都放在开头。嗯。对领导们得先发言，得先站到台上挨个儿过一遍。啊，是。他还介绍一下都是谁。那是不，那是不是太枯燥了一点。先得介先得先得介绍一下，就是走流程嘛，哎呀。就主持人肯定要介绍一下是哪个领导。那时间得缩短点，不然那家伙那大伙都睡着了。嗯嗯。对就赶紧介绍，然后赶紧就开始正式的一些活动。对啊，那肯定的。对，然后。可以可以。可以可以。那学生肯定辅导员也得来呗，不然谁管得住他们啊。导员导员可以坐在学生，学生，每个班级的旁边，对是可以，是可以这个样子的。对，可以站着管理一下纪律什么的。啊一个月一个月就是分开做，然后辅导员坐边上。对对对。对，对对对，哎，你们都你们学校都是这个样子的吗？对嗯。唉，都差不多吧，应该。对，对。嗯。嗯，哦我说你们的高中啊，高中都是这个样子的，对吧？对对，然后。辅导员，领导，什么的都好了。就是人人有点不够吧。那在请点儿学生家长，但又不能都请都请人忒多了。嗯可以可以。嗯。对也没地儿没，也没人坐。那可以。做点儿邀请函，做点儿邀请函给每个班级同学发过去，然后看看谁谁家长能来谁就来。对。谁需要，对，每个班级发一些就好了。每个班得还得固定一下人数。而我觉得这个邀请函一定要做的比较有新意一点儿，就是做的比较能吸引大家。对每个班发发一点。做个海报算了，做邀请函耗时耗时<EFBFBD>

    def run_pipeline(self, model_id: str, documents: str) -> Dict[str, Any]:
        p = pipeline(task=self.task, model=model_id)
        result = p(documents=documents)
        return result

    @unittest.skipUnless(test_level() >= 2, 'skip test in current test level')
    def test_run_with_doc(self):
        logger.info(
            'Run doc extractive summarization (PoNet) with one document ...')

        result = self.run_pipeline(
            model_id=self.ponet_doc_model_id, documents=self.sentences)
        print(result[OutputKeys.TEXT])

    @unittest.skipUnless(test_level() >= 2, 'skip test in current test level')
    def test_run_with_topic(self):
        logger.info(
            'Run topic extractive summarization (PoNet) with one document ...')

        result = self.run_pipeline(
            model_id=self.ponet_topic_model_id, documents=self.sentences)
        print(result[OutputKeys.TEXT])


if __name__ == '__main__':
    unittest.main()