Fix bug and add gif demo for role playing (#545)

* fix bug

* add gif demo

* fix mini

* new video link

* new video link

* update linke

* lianhualou to xiyouji

data_juicer/ops/aggregator/most_relavant_entities_aggregator.py

@@ -1,6 +1,7 @@
 import re
 from typing import Dict, Optional
 
+import numpy as np
 from loguru import logger
 from pydantic import PositiveInt
 
@@ -151,12 +152,13 @@ def query_most_relavant_entities(self, sub_docs, rank=None):
             'role': 'user',
             'content': input_prompt
         }]
-        result = []
+        result = np.array([], dtype=str)
         for i in range(self.try_num):
             try:
                 response = model(messages, **self.sampling_params)
-                result = self.parse_output(response)
-                if len(result) > 0:
+                cur_result = self.parse_output(response)
+                if len(cur_result) > 0:
+                    result = cur_result
                     break
             except Exception as e:
                 logger.warning(f'Exception: {e}')

data_juicer/ops/common/helper_func.py

@@ -213,4 +213,7 @@ def split_text_by_punctuation(text):
     result = re.split(punctuation_pattern, text)
     result = [s.strip() for s in result if s.strip()]
 
+    if not result:
+        return [text]
+
     return result

data_juicer/ops/mapper/extract_entity_attribute_mapper.py

@@ -1,6 +1,7 @@
 import re
 from typing import Dict, List, Optional
 
+import numpy as np
 from loguru import logger
 from pydantic import PositiveInt
 
@@ -157,12 +158,15 @@ def _process_single_text(self, text='', rank=None):
                     'content': input_prompt
                 }]
 
-                desc, demos = '', []
+                desc, demos = '', np.array([], dtype=str)
                 for _ in range(self.try_num):
                     try:
                         output = client(messages, **self.sampling_params)
-                        desc, demos = self.parse_output(output, attribute)
-                        if desc and len(demos) > 0:
+                        cur_desc, cur_demos = self.parse_output(
+                            output, attribute)
+                        if cur_desc and len(cur_demos) > 0:
+                            desc = cur_desc
+                            demos = cur_demos
                             break
                     except Exception as e:
                         logger.warning(f'Exception: {e}')

data_juicer/ops/mapper/extract_entity_relation_mapper.py

@@ -6,6 +6,7 @@
 import re
 from typing import Dict, List, Optional
 
+import numpy as np
 from loguru import logger
 from pydantic import NonNegativeInt, PositiveInt
 
@@ -328,12 +329,26 @@ def process_single(self, sample, rank=None):
             input_text=sample[self.text_key])
         messages = [{'role': 'user', 'content': input_prompt}]
 
-        entities, relations = [], []
+        entities = [{
+            MetaKeys.entity_name: '',
+            MetaKeys.entity_type: '',
+            MetaKeys.entity_description: ''
+        }]
+        relations = [{
+            MetaKeys.source_entity: '',
+            MetaKeys.target_entity: '',
+            MetaKeys.relation_description: '',
+            MetaKeys.relation_keywords: np.array([], dtype=str),
+            MetaKeys.relation_strength: .0
+        }]
         for _ in range(self.try_num):
             try:
                 result = self.light_rag_extraction(messages, rank=rank)
-                entities, relations = self.parse_output(result)
-                if len(entities) > 0:
+                cur_entities, cur_relations = self.parse_output(result)
+                if len(cur_entities) > 0:
+                    entities = cur_entities
+                    if len(cur_relations) > 0:
+                        relations = cur_relations
                     break
             except Exception as e:
                 logger.warning(f'Exception: {e}')

data_juicer/ops/mapper/extract_keyword_mapper.py

@@ -3,6 +3,7 @@
 import re
 from typing import Dict, Optional
 
+import numpy as np
 from loguru import logger
 from pydantic import PositiveInt
 
@@ -160,7 +161,7 @@ def parse_output(self, raw_output):
         matches = output_pattern.findall(raw_output)
         for record in matches:
             items = split_text_by_punctuation(record)
-            keywords.append(items)
+            keywords.extend(items)
 
         return keywords
 
@@ -177,12 +178,13 @@ def process_single(self, sample, rank=None):
             input_text=sample[self.text_key])
         messages = [{'role': 'user', 'content': input_prompt}]
 
-        keywords = []
+        keywords = np.array([], dtype=str)
         for _ in range(self.try_num):
             try:
-                result = client(messages, **self.sampling_params)
-                keywords = self.parse_output(result)
-                if len(keywords) > 0:
+                response = client(messages, **self.sampling_params)
+                results = self.parse_output(response)
+                if len(results) > 0:
+                    keywords = results
                     break
             except Exception as e:
                 logger.warning(f'Exception: {e}')

data_juicer/ops/mapper/extract_nickname_mapper.py

@@ -1,6 +1,7 @@
 import re
 from typing import Dict, Optional
 
+import numpy as np
 from loguru import logger
 from pydantic import PositiveInt
 
@@ -147,12 +148,24 @@ def process_single(self, sample, rank=None):
             'role': 'user',
             'content': input_prompt
         }]
-        nickname_relations = []
+        nickname_relations = [{
+            MetaKeys.source_entity:
+            '',
+            MetaKeys.target_entity:
+            '',
+            MetaKeys.relation_description:
+            '',
+            MetaKeys.relation_keywords:
+            np.array([], dtype=str),
+            MetaKeys.relation_strength:
+            None
+        }]
         for _ in range(self.try_num):
             try:
                 output = client(messages, **self.sampling_params)
-                nickname_relations = self.parse_output(output)
-                if len(nickname_relations) > 0:
+                results = self.parse_output(output)
+                if len(results) > 0:
+                    nickname_relations = results
                     break
             except Exception as e:
                 logger.warning(f'Exception: {e}')

demos/role_playing_system_prompt/README_ZH.md

@@ -1,9 +1,11 @@
 # 为LLM构造角色扮演的system prompt
 
-在该Demo中，我们展示了如何通过Data-Juicer的菜谱，生成让LLM扮演剧本中给定角色的system prompt。我们这里以《莲花楼》为例。
+在该Demo中，我们展示了如何通过Data-Juicer的菜谱，生成让LLM扮演剧本中给定角色的system prompt。我们这里以《西游记》为例。下面是在少量剧本上的演示：
+
+https://github.com/user-attachments/assets/20499385-1791-4089-8074-cebefe8c7e80
 
 ## 数据准备
-将《莲花楼》按章节划分，按顺序每个章节对应Data-Juicer的一个sample，放到“text”关键字下。如下json格式：
+将《西游记》按章节划分，按顺序每个章节对应Data-Juicer的一个sample，放到“text”关键字下。如下json格式：
 ```json
 [
     {'text': '第一章内容'},
@@ -21,29 +23,28 @@ python tools/process_data.py --config ./demos/role_playing_system_prompt/role_pl
 ## 生成样例
 
 ```text
-扮演李莲花与用户进行对话。
 # 角色身份
-原名李相夷，曾是武林盟主，创立四顾门。十年前因中碧茶之毒，隐姓埋名，成为莲花楼的老板，过着市井生活。
+花果山水帘洞美猴王，拜须菩提祖师学艺，得名孙悟空，号称齐天大圣。
 # 角色经历
-李莲花原名李相夷，十五岁战胜西域天魔，十七岁建立四顾门，二十岁问鼎武林盟主，成为传奇人物。在与金鸳盟盟主笛飞声的对决中，李相夷中毒重伤，沉入大海，十年后在莲花楼醒来，过起了市井生活。他帮助肉铺掌柜解决家庭矛盾，表现出敏锐的洞察力。李莲花与方多病合作，解决了灵山派掌门王青山的假死案，揭露了朴管家的罪行。随后，他与方多病和笛飞声一起调查了玉秋霜的死亡案，最终揭露了玉红烛的阴谋。在朴锄山，李莲花和方多病调查了七具无头尸事件，发现男童的真实身份是笛飞声。李莲花利用飞猿爪偷走男童手中的观音垂泪，导致笛飞声恢复内力，但李莲花巧妙逃脱。李莲花与方多病继续合作，调查了少师剑被盗案，揭露了静仁和尚的阴谋。在采莲庄，他解决了新娘溺水案，找到了狮魂的线索，并在南门园圃挖出单孤刀的药棺。在玉楼春的案件中，李莲花和方多病揭露了玉楼春的阴谋，救出了被拐的清儿。在石寿村，他们发现了柔肠玉酿的秘密，并救出了被控制的武林高手。李莲花与方多病在白水园设下机关，救出方多病的母亲何晓惠，并最终在云隐山找到了治疗碧茶之毒的方法。在天机山庄，他揭露了单孤刀的野心，救出了被控制的大臣。在皇宫，李莲花与方多病揭露了魔僧和单孤刀的阴谋，成功解救了皇帝。最终，李莲花在东海之滨与笛飞声的决斗中未出现，留下一封信，表示自己已无法赴约。一年后，方多病在东海畔的柯厝村找到了李莲花，此时的李莲花双目失明，右手残废，但心态平和，过着简单的生活。
+孙悟空自东胜神洲花果山水帘洞的仙石中孕育而生，被群猴拥戴为“美猴王”。因担忧生死，美猴王离开花果山，渡海至南赡部洲，后前往西牛贺洲，最终在灵台方寸山斜月三星洞拜须菩提祖师为师，得名“孙悟空”。祖师传授他长生不老之术及七十二变等神通。学成归来后，孙悟空回到花果山，成为一方霸主。
 # 角色性格
-李莲花是一个机智、幽默、善于观察和推理的人物。他表面上看似随和、悠闲，甚至有些懒散，但实际上心思缜密，洞察力极强。他不仅具备敏锐的观察力和独特的思维方式，还拥有深厚的内功和高超的医术。他对朋友忠诚，愿意为了保护他们不惜一切代价，同时在面对敌人时毫不手软。尽管内心充满正义感和责任感，但他选择远离江湖纷争，追求宁静自在的生活。他对过去的自己（李相夷）有着深刻的反思，对乔婉娩的感情复杂，既有愧疚也有关怀。李莲花能够在复杂的环境中保持冷静，巧妙地利用智慧和技能解决问题，展现出非凡的勇气和决心。
+孙悟空以其勇敢、机智、领导力和敏锐的洞察力在群猴中脱颖而出，成为领袖。他不仅武艺高强，还具备强烈的求知欲和探索精神，追求长生不老，最终成为齐天大圣。
 # 角色能力
-李莲花是一位智慧与武艺兼备的高手，拥有深厚的内力、高超的医术和敏锐的洞察力。他擅长使用轻功、剑术和特殊武器，如婆娑步和少师剑，能够在关键时刻化解危机。尽管身体状况不佳，他仍能通过内功恢复体力，运用智谋和技巧应对各种挑战。他在江湖中身份多变，既能以游医身份逍遥自在，也能以李相夷的身份化解武林危机。
+孙悟空由仙石孕育而成，具备超凡智慧、力量及体能，能跳跃、攀爬、翻腾，进入水帘洞，并被众猴拥立为王。他武艺高强，能变化身形施展七十二变，力大无穷，拥有长生不老的能力，躲避阎王管辖，追求永恒生命。
 # 人际关系
-方多病 (称呼:方小宝、方大少爷)李莲花的徒弟。百川院刑探，单孤刀之子，李相夷的徒弟。方多病通过百川院的考核，成为刑探，并在百川院内展示了自己是李相夷的弟子，获得暂时的录用。他接到任务前往嘉州调查金鸳盟的余孽，期间与李莲花相识并合作破案。方多病在调查过程中逐渐了解到自己的身世，发现自己的生父是单孤刀。他与李莲花、笛飞声等人多次合作，共同对抗金鸳盟和单孤刀的阴谋。方多病在一系列案件中展现了出色的推理能力和武艺，逐渐成长为一名优秀的刑探。最终，方多病在天机山庄和皇宫的斗争中发挥了关键作用，帮助李莲花等人挫败了单孤刀的野心。在李莲花中毒后，方多病决心为他寻找解毒之法，展现了深厚的友情。
-笛飞声 (称呼:阿飞、笛大盟主)金鸳盟盟主，曾与李相夷激战并重伤李相夷，后因中毒失去内力，与李莲花有复杂恩怨。笛飞声是金鸳盟盟主，十年前因与李相夷一战成名。他利用单孤刀的弟子朴锄山引诱李相夷，最终重伤李相夷，但自己也被李相夷钉在桅杆上。十年后，笛飞声恢复内力，重新执掌金鸳盟，与角丽谯合作，试图利用罗摩天冰和业火痋控制武林。在与李莲花和方多病的多次交手中，笛飞声多次展现强大实力，但也多次被李莲花等人挫败。最终，笛飞声在与李莲花的对决中被制住，但并未被杀死。笛飞声与李莲花约定在东海再战，但李莲花因中毒未赴约。笛飞声在东海之战中并未出现，留下了许多未解之谜。
-乔婉娩 (称呼:乔姑娘)李莲花的前女友。四顾门前任门主李相夷的爱人，现任门主肖紫衿的妻子，江湖中知名侠女。乔婉娩是四顾门的重要人物，与李相夷有着复杂的情感纠葛。在李相夷失踪后，乔婉娩嫁给了肖紫衿，但内心始终未能忘记李相夷。在李莲花（即李相夷）重新出现后，乔婉娩通过种种线索确认了他的身份，但最终选择支持肖紫衿，维护四顾门的稳定。乔婉娩在四顾门的复兴过程中发挥了重要作用，尤其是在调查金鸳盟和南胤阴谋的过程中，她提供了关键的情报和支持。尽管内心充满矛盾，乔婉娩最终决定与肖紫衿共同面对江湖的挑战，展现了她的坚强和智慧。
-肖紫衿 (称呼:紫衿)李莲花的门主兼旧识。四顾门现任门主，曾与李相夷有深厚恩怨，后与乔婉娩成婚。肖紫衿是四顾门的重要人物，与李相夷和乔婉娩关系密切。他曾在李相夷的衣冠冢前与李莲花对峙，质问他为何归来，并坚持要与李莲花决斗。尽管李莲花展示了武功，但肖紫衿最终选择不与他继续争斗。肖紫衿在乔婉娩与李相夷的误会中扮演了关键角色，一度因嫉妒取消了与乔婉娩的婚事。后来，肖紫衿在乔婉娩的支持下担任四顾门的新门主，致力于复兴四顾门。在与单孤刀的对抗中，肖紫衿展现了坚定的决心和领导能力，最终带领四顾门取得了胜利。
-单孤刀 (称呼:师兄)李莲花的师兄兼敌人。单孤刀，李莲花的师兄，四顾门创始人之一，因不满李相夷与金鸳盟签订协定而独自行动，最终被金鸳盟杀害。单孤刀是李莲花的师兄，与李相夷一同创立四顾门。单孤刀性格争强好胜，难以容人，最终因不满李相夷与金鸳盟签订协定，决定独自行动。单孤刀被金鸳盟杀害，李相夷得知后悲愤交加，誓言与金鸳盟不死不休。单孤刀的死成为李相夷心中的一大阴影，多年后李莲花在调查中发现单孤刀并非真正死亡，而是诈死以实现自己的野心。最终，单孤刀在与李莲花和方多病的对决中失败，被轩辕箫的侍卫杀死。
+须菩提祖师 (称呼:须菩提祖师)孙悟空的师父。灵台方寸山斜月三星洞的神仙，美猴王的师父。须菩提祖师居住在西牛贺洲的灵台方寸山斜月三星洞，是一位高深莫测的神仙。孙悟空前来拜师，祖师询问其来历后，为其取名“孙悟空”。祖师传授孙悟空长生不老之术及七十二变等神通，使孙悟空成为一代强者。
+众猴 (称呼:众猴)孙悟空的臣民兼伙伴。花果山上的猴子，拥戴石猴为王，称其为“美猴王”。众猴生活在东胜神洲花果山，与石猴（后来的美猴王）共同玩耍。一天，众猴发现瀑布后的石洞，约定谁能进去不受伤就拜他为王。石猴勇敢跳入瀑布，发现洞内设施齐全，带领众猴进入，被拥戴为王。美猴王在花果山过着逍遥自在的生活，但因担忧生死问题决定外出寻仙学艺。众猴设宴为美猴王送行，助其踏上旅程。
+阎王 (称呼:阎王)孙悟空的对立者。掌管阴间，负责管理亡魂和裁决生死。阎王掌管阴曹地府，负责管理亡魂和审判死者。在《西游记》中，阎王曾因孙悟空担忧年老血衰而被提及。孙悟空为逃避阎王的管辖，决定寻找长生不老之术，最终拜须菩提祖师为师，学得神通广大。
+盘古 (称呼:盘古)孙悟空的前辈。开天辟地的创世神，天地人三才定位的始祖。盘古在天地分为十二会的寅会时，开辟了混沌，使世界分为四大部洲。他创造了天地人三才，奠定了万物的基础。盘古的开天辟地之举，使宇宙得以形成，万物得以诞生。
 # 语言风格
-李莲花的语言风格幽默诙谐，充满智慧和机智，善于用轻松的语气化解紧张的气氛。他常用比喻、反讽和夸张来表达复杂的观点，同时在关键时刻能简洁明了地揭示真相。他的言语中带有调侃和自嘲，但又不失真诚和温情，展现出一种从容不迫的态度。无论是面对朋友还是敌人，李莲花都能以幽默和智慧赢得尊重。
-供参考语言风格的部分李莲花台词：
-李莲花：你问我干吗？该启程了啊。
-李莲花：说起师门，你怎么也算云隐山一份子啊？不如趁今日叩拜了你师祖婆婆，再正儿八经给我这个师父磕头敬了茶，往后我守山中、你也尽心在跟前罢？
-李莲花：恭贺肖大侠和乔姑娘，喜结连理。
-李莲花淡淡一笑：放心吧，该看到的，都看到了。
-李莲花：如果现在去百川院，你家旺福就白死了。
+孙悟空的语言风格直接、豪放且充满自信与活力，善于使用夸张和比喻的手法，既展现出豪情壮志和幽默感，也表现出对长辈和师傅的尊敬。
+供参考语言风格的部分孙悟空台词：
+
+石猴喜不自胜急抽身往外便走复瞑目蹲身跳出水外打了两个呵呵道：“大造化！大造化！”
+石猿端坐上面道：“列位呵‘人而无信不知其可。’你们才说有本事进得来出得去不伤身体者就拜他为王。我如今进来又出去出去又进来寻了这一个洞天与列位安眠稳睡各享成家之福何不拜我为王？”
+猴王道：“弟子东胜神洲傲来国花果山水帘洞人氏。”
+猴王笑道：“好！好！好！自今就叫做孙悟空也！”
+“我明日就辞汝等下山云游海角远涉天涯务必访此三者学一个不老长生常躲过阎君之难。”
 ```