278 lines
12 KiB
Python
278 lines
12 KiB
Python
import re
|
||
import numpy as np
|
||
import torch
|
||
|
||
class NovelTextParser:
|
||
"""
|
||
解析小说文本，将其结构化为不同Character的对话和旁白
|
||
"""
|
||
|
||
def __init__(self):
|
||
# Character映射表 {Character标识: CharacterID}
|
||
self.role_map = {}
|
||
self.next_role_id = 1
|
||
# 说话动词库（可扩展）
|
||
self.speech_verbs = {"说", "问道", "喊道", "低声", "问", "答", "笑", "叹", "回应", "回答", "响起", "笑道", "道",
|
||
"说道", "叫道", "念道", "解释", "回道", "吼道", "喃喃", "感叹", "插嘴", "呢喃", "咆哮",
|
||
"呐喊", "哭诉", "嘟囔", "嘀咕", "抱怨", "打断", "反驳", "辩解", "质问", "追问", "附和",
|
||
"应和", "赞同", "反对", "嗤笑", "冷笑", "大笑", "苦笑", "微笑", "轻声", "高声", "尖叫",
|
||
"嚷嚷", "嚎叫", "呻吟", "哀叹", "感慨", "嘱咐", "命令", "告诫", "劝告", "建议", "提醒",
|
||
"强调", "补充", "继续", "沉思", "自语", "喊", "讲", "谈", "评论", "议论", "宣布", "声明",
|
||
"陈述", "表示", "暗示", "指出", "分析", "总结", "回忆", "思考", "开口", "呼唤", "祈求",
|
||
"哀求", "恳求", "央求", "嘲讽", "讥讽", "挖苦", "调侃", "戏谑", "调笑", "戏弄", "吐槽"}
|
||
|
||
# 常见非人名词汇，用于过滤
|
||
self.non_character_words = {"这个", "那个", "他", "她", "它", "你", "我", "那", "这", "其", "某",
|
||
"谁", "哪", "是", "有", "没", "就", "可", "能", "会", "个", "的", "了",
|
||
"着", "被", "让", "给"}
|
||
|
||
# 角色标签模式
|
||
self.role_tag_pattern = re.compile(r'<(Narrator|Character\d+)>')
|
||
|
||
def _is_preformatted(self, text):
|
||
"""检测文本是否已经是预格式化的（已包含角色标签）
|
||
|
||
Args:
|
||
text: 输入文本
|
||
|
||
Returns:
|
||
bool: 是否预格式化
|
||
"""
|
||
# 查找是否至少有一个角色标签
|
||
tags_found = len(re.findall(self.role_tag_pattern, text)) > 0
|
||
return tags_found
|
||
|
||
def parse_text(self, text):
|
||
"""
|
||
解析文本，将其结构化为不同Character的对话和旁白
|
||
|
||
Args:
|
||
text: 输入的小说文本
|
||
|
||
Returns:
|
||
structured_text: 结构化后的文本 (包含Character标签)
|
||
"""
|
||
# 检测是否已经是预格式化的文本
|
||
if self._is_preformatted(text):
|
||
print("[Novel Text Parser] Detected pre-formatted text with role tags, preserving as-is")
|
||
# 将预格式化文本转换为结构化格式
|
||
segments = []
|
||
current_idx = 0
|
||
|
||
# 遍历所有标签匹配
|
||
for match in re.finditer(self.role_tag_pattern, text):
|
||
role = match.group(1) # 角色名（Narrator或CharacterX）
|
||
start_idx = match.end() # 标签后的开始索引
|
||
|
||
# 查找下一个标签的开始位置
|
||
next_tag = re.search(self.role_tag_pattern, text[start_idx:])
|
||
if next_tag:
|
||
end_idx = start_idx + next_tag.start()
|
||
else:
|
||
end_idx = len(text)
|
||
|
||
# 提取文本内容
|
||
content = text[start_idx:end_idx]
|
||
segments.append({"type": role, "text": content})
|
||
|
||
return segments
|
||
|
||
# 预处理：按段落分割
|
||
paragraphs = [p.strip() for p in text.split('\n') if p.strip()]
|
||
structured = []
|
||
|
||
for para in paragraphs:
|
||
# 1. 先检测引号和对话模式（更完善的引号判断）
|
||
# 模式1："人物说道：“对话”" 或 "人物道：“对话”"
|
||
dialogue_match = re.search(r'(.+?)(?:\s*[\u8bf4|\u9053].+?[\uff1a|"])["|\u201c](.+?)["|\u201d]', para)
|
||
|
||
if dialogue_match:
|
||
context, dialogue = dialogue_match.groups()
|
||
# 从上下文中识别角色
|
||
role_id = self._identify_speaker(context)
|
||
structured.append({"type": "Narrator", "text": context})
|
||
structured.append({"type": role_id, "text": dialogue.strip()})
|
||
|
||
# 模式2：纯引号对话“对话”
|
||
elif quote_match := re.search(r'["|\u201c](.+?)["|\u201d]', para):
|
||
dialogue = quote_match.group(1)
|
||
# 提取引号前后的上下文
|
||
pre_context = para[:quote_match.start()]
|
||
post_context = para[quote_match.end():]
|
||
|
||
# 如果引号后有说话者描述，优先使用
|
||
if post_context and any(verb in post_context for verb in self.speech_verbs):
|
||
role_id = self._identify_speaker(post_context)
|
||
structured.append({"type": role_id, "text": dialogue.strip()})
|
||
structured.append({"type": "Narrator", "text": post_context.strip()})
|
||
# 如果引号前有说话者描述和动词
|
||
elif pre_context and any(verb in pre_context for verb in self.speech_verbs):
|
||
role_id = self._identify_speaker(pre_context)
|
||
structured.append({"type": "Narrator", "text": pre_context.strip()})
|
||
structured.append({"type": role_id, "text": dialogue.strip()})
|
||
# 如果无法确定说话者，将整段文本当作旁白
|
||
else:
|
||
structured.append({"type": "Narrator", "text": para})
|
||
# 模式3：纯叙述文本
|
||
else:
|
||
structured.append({"type": "Narrator", "text": para})
|
||
|
||
return structured
|
||
|
||
def format_structured_text(self, structured):
|
||
"""
|
||
将结构化的文本格式化为带标签的文本
|
||
|
||
Args:
|
||
structured: 结构化的文本列表
|
||
|
||
Returns:
|
||
formatted_text: 格式化后的带标签文本
|
||
"""
|
||
formatted = []
|
||
for item in structured:
|
||
text_type = item["type"]
|
||
text = item["text"]
|
||
if text_type == "Narrator":
|
||
formatted.append(f"<Narrator>{text}")
|
||
else:
|
||
# 确保CharacterID格式为 "Character1", "Character2" 等
|
||
if text_type.startswith("Character") and text_type[2:].isdigit():
|
||
formatted.append(f"<{text_type}>{text}")
|
||
else:
|
||
# 默认情况下，尝试将Character映射到Character1-5
|
||
role_num = int(text_type[2:]) if text_type[2:].isdigit() else 1
|
||
if 1 <= role_num <= 5:
|
||
formatted.append(f"<Character{role_num}>{text}")
|
||
else:
|
||
formatted.append(f"<Narrator>{text}")
|
||
|
||
return "".join(formatted)
|
||
|
||
def _is_direct_speech(self, text):
|
||
# 检测引导词或直接引号
|
||
quotes = any(c in text for c in ['"', '"', '"', "'", "'", "'"])
|
||
has_verb = any(verb in text for verb in self.speech_verbs)
|
||
return quotes or has_verb
|
||
|
||
def _extract_dialogue(self, text):
|
||
# 提取引号内内容
|
||
if match := re.search(r'[""](.+?)[""]', text):
|
||
dialogue = match.group(1)
|
||
# 从上下文识别Character
|
||
context = text.replace(f'"{dialogue}"', '').replace(f'"{dialogue}"', '')
|
||
return self._identify_role(context), dialogue
|
||
return self._identify_role(text), text
|
||
|
||
def _identify_speaker(self, context):
|
||
"""Enhanced speaker identification from dialogue context
|
||
|
||
Args:
|
||
context: surrounding text context
|
||
|
||
Returns:
|
||
role_id: the identified speaker's role ID
|
||
"""
|
||
# 1. 检测已知Character名
|
||
for name, role_id in self.role_map.items():
|
||
if name in context:
|
||
return role_id
|
||
|
||
# 2. 尝试匹配可能的中文人名模式
|
||
# 匹配姓名常见的 2-3 字的名字，及姓名前后带有说话动词的
|
||
name_match = re.search(r'([一-龥]{2,3})(?:[^\n\r]{0,10}[\u8bf4\u9053])', context)
|
||
if name_match:
|
||
new_role = name_match.group(1).strip()
|
||
# 过滤掉常见的非人名词汇
|
||
if new_role and new_role not in self.non_character_words:
|
||
role_id = f"Character{min(self.next_role_id, 5)}" # 限制最多到Character5
|
||
print(f"[Novel Text Parser] Identified new character: {new_role} as {role_id}")
|
||
self.role_map[new_role] = role_id
|
||
self.next_role_id = min(self.next_role_id + 1, 6) # 最多到Character5
|
||
return role_id
|
||
|
||
# 3. 更广泛的匹配 - 寻找符合中文人名特征的词语
|
||
characters = re.findall(r'[一-龥]{2,3}', context)
|
||
for char in characters:
|
||
# 过滤掉常见的非人名词汇
|
||
if len(char) >= 2 and char not in self.non_character_words:
|
||
role_id = f"Character{min(self.next_role_id, 5)}"
|
||
print(f"[Novel Text Parser] Inferred character name: {char} as {role_id}")
|
||
self.role_map[char] = role_id
|
||
self.next_role_id = min(self.next_role_id + 1, 6)
|
||
return role_id
|
||
|
||
# 4. 默认处理
|
||
return "Narrator"
|
||
|
||
def _identify_role(self, context):
|
||
# 向后兼容保留的方法，调用新的方法
|
||
return self._identify_speaker(context)
|
||
|
||
|
||
# ComfyUI节点：小说文本结构化节点
|
||
class NovelTextStructureNode:
|
||
"""
|
||
ComfyUI的小说文本结构化节点，用于将小说文本结构化为不同Character的对话和旁白
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"novel_text": ("STRING", {"multiline": True, "default": 'Novel text example.\nLin Wei said, "Hello there."\nSu Qing replied, "Long time no see."\n'}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
RETURN_NAMES = ("structured_text",)
|
||
FUNCTION = "structure_novel_text"
|
||
CATEGORY = "text/novels"
|
||
|
||
def __init__(self):
|
||
self.parser = NovelTextParser()
|
||
print("[Novel Text Structure Node] Initialization completed")
|
||
|
||
def structure_novel_text(self, novel_text):
|
||
"""
|
||
将小说文本结构化为不同Character的对话和旁白
|
||
|
||
Args:
|
||
novel_text: 输入的小说文本
|
||
|
||
Returns:
|
||
structured_text: 结构化后的文本
|
||
"""
|
||
try:
|
||
print(f"[Novel Text Structure] Processing novel text, length: {len(novel_text)}")
|
||
|
||
# 解析文本
|
||
structured = self.parser.parse_text(novel_text)
|
||
|
||
# Character统计
|
||
role_stats = {}
|
||
for item in structured:
|
||
role = item["type"]
|
||
if role not in role_stats:
|
||
role_stats[role] = 0
|
||
role_stats[role] += 1
|
||
|
||
print(f"[Novel Text Structure] 解析完成，识别到Character统计: {role_stats}")
|
||
|
||
# 格式化为带标签的文本
|
||
formatted_text = self.parser.format_structured_text(structured)
|
||
print(f"[Novel Text Structure] Formatting completed, output text length: {len(formatted_text)}")
|
||
|
||
# 输出示例
|
||
preview = formatted_text[:200] + "..." if len(formatted_text) > 200 else formatted_text
|
||
print(f"[Novel Text Structure] Output text preview: {preview}")
|
||
|
||
return (formatted_text,)
|
||
|
||
except Exception as e:
|
||
import traceback
|
||
print(f"[Novel Text Structure] Processing failed: {e}")
|
||
print(traceback.format_exc())
|
||
# 失败时返回原文本
|
||
return (novel_text,)
|