test_most_relavant_entities_aggregator.py

import unittest

from loguru import logger

from data_juicer.core.data import NestedDataset as Dataset
from data_juicer.ops.aggregator import MostRelavantEntitiesAggregator
from data_juicer.utils.unittest_utils import DataJuicerTestCaseBase, SKIPPED_TESTS


@SKIPPED_TESTS.register_module()
class MostRelavantEntitiesAggregatorTest(DataJuicerTestCaseBase):

    def _run_helper(self, op, samples):

        # before runing this test, set below environment variables:
        # export OPENAI_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1/
        # export OPENAI_API_KEY=your_dashscope_key

        dataset = Dataset.from_list(samples)
        new_dataset = op.run(dataset)

        for data in new_dataset:
            for k in data:
                logger.info(f"{k}: {data[k]}")

        self.assertEqual(len(new_dataset), len(samples))

    def test_default_aggregator(self):
        samples = [
            {
                'text': [
                    "十年前，李相夷十五岁战胜西域天魔成为天下第一高手，十七岁建立四顾门，二十岁问鼎武林盟主，成为传奇人物。",
                    "有人视李相夷为中原武林的希望，但也有人以战胜他为目标，包括魔教金鸳盟盟主笛飞声。笛飞声设计加害李相夷的师兄单孤刀，引得李相夷与之一战。",
                    '在东海的一艘船上，李相夷独自一人对抗金鸳盟的高手，最终击败了大部分敌人。笛飞声突然出现，两人激战，李相夷在战斗中中毒，最终被笛飞声重伤，船只爆炸，李相夷沉入大海。',
                    '十年后，李莲花在一个寒酸的莲花楼内醒来，表现出与李相夷截然不同的性格。他以神医的身份在小镇上行医，但生活贫困。',
                    '小镇上的皮影戏摊讲述李相夷和笛飞声的故事，孩子们争论谁赢了。风火堂管事带着人来找李莲花，要求他救治一个“死人”。'
                ]
            },
        ]
        
        op = MostRelavantEntitiesAggregator(
            api_model='qwen2.5-72b-instruct',
            entity='李莲花',
            query_entity_type='人物'
        )
        self._run_helper(op, samples)
    
    def test_input_output(self):
        samples = [
            {
                'dj_result':{
                    'events': [
                        "十年前，李相夷十五岁战胜西域天魔成为天下第一高手，十七岁建立四顾门，二十岁问鼎武林盟主，成为传奇人物。",
                        "有人视李相夷为中原武林的希望，但也有人以战胜他为目标，包括魔教金鸳盟盟主笛飞声。笛飞声设计加害李相夷的师兄单孤刀，引得李相夷与之一战。",
                        '在东海的一艘船上，李相夷独自一人对抗金鸳盟的高手，最终击败了大部分敌人。笛飞声突然出现，两人激战，李相夷在战斗中中毒，最终被笛飞声重伤，船只爆炸，李相夷沉入大海。',
                        '十年后，李莲花在一个寒酸的莲花楼内醒来，表现出与李相夷截然不同的性格。他以神医的身份在小镇上行医，但生活贫困。',
                        '小镇上的皮影戏摊讲述李相夷和笛飞声的故事，孩子们争论谁赢了。风火堂管事带着人来找李莲花，要求他救治一个“死人”。'
                    ]
                }
            },
        ]

        op = MostRelavantEntitiesAggregator(
            api_model='qwen2.5-72b-instruct',
            entity='李莲花',
            query_entity_type='人物',
            input_key='dj_result.events',
            output_key='dj_result.relavant_roles'
        )
        self._run_helper(op, samples)

    def test_max_token_num(self):
        samples = [
            {
                'text': [
                    "十年前，李相夷十五岁战胜西域天魔成为天下第一高手，十七岁建立四顾门，二十岁问鼎武林盟主，成为传奇人物。",
                    "有人视李相夷为中原武林的希望，但也有人以战胜他为目标，包括魔教金鸳盟盟主笛飞声。笛飞声设计加害李相夷的师兄单孤刀，引得李相夷与之一战。",
                    '在东海的一艘船上，李相夷独自一人对抗金鸳盟的高手，最终击败了大部分敌人。笛飞声突然出现，两人激战，李相夷在战斗中中毒，最终被笛飞声重伤，船只爆炸，李相夷沉入大海。',
                    '十年后，李莲花在一个寒酸的莲花楼内醒来，表现出与李相夷截然不同的性格。他以神医的身份在小镇上行医，但生活贫困。',
                    '小镇上的皮影戏摊讲述李相夷和笛飞声的故事，孩子们争论谁赢了。风火堂管事带着人来找李莲花，要求他救治一个“死人”。'
                ]
            },
        ]
        op = MostRelavantEntitiesAggregator(
            api_model='qwen2.5-72b-instruct',
            entity='李莲花',
            query_entity_type='人物',
            max_token_num=40
        )
        self._run_helper(op, samples)

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