237 lines
8.8 KiB
Python
237 lines
8.8 KiB
Python
"""
|
||
关系抽取模块
|
||
使用多个LLM（GPT-4o, Doubao-pro, GLM-4）并行进行关系抽取
|
||
"""
|
||
|
||
import json
|
||
import logging
|
||
from typing import List, Dict, Any, Optional
|
||
|
||
from ..prompts.re_prompts import REPromptBuilder
|
||
from ..utils.llm_client import LLMClient, LLMProvider
|
||
from ..utils.config import Config
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class REExtractor:
|
||
"""关系抽取器"""
|
||
|
||
def __init__(
|
||
self,
|
||
models: Optional[List[str]] = None,
|
||
config: Optional[Config] = None,
|
||
):
|
||
"""
|
||
初始化RE提取器
|
||
|
||
Args:
|
||
models: 使用的模型列表，默认使用多个模型并行
|
||
config: 配置对象
|
||
"""
|
||
self.config = config or Config()
|
||
|
||
# 默认使用多个模型（优先使用硅基流动）
|
||
if models is None:
|
||
models = []
|
||
# 优先使用硅基流动（使用两个模型）
|
||
if self.config.SILICONFLOW_API_KEY:
|
||
models.append(("siliconflow", "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"))
|
||
models.append(("siliconflow", "Qwen/Qwen2.5-7B-Instruct"))
|
||
# 如果有其他API密钥，也添加
|
||
if self.config.OPENAI_API_KEY:
|
||
models.append(("openai", "gpt-4o"))
|
||
if self.config.VOLCENGINE_ACCESS_KEY and self.config.VOLCENGINE_SECRET_KEY:
|
||
models.append(("doubao", "doubao-pro-32k"))
|
||
if self.config.ZHIPUAI_API_KEY:
|
||
models.append(("glm", "glm-4-airx"))
|
||
|
||
# 如果都没有配置，至少使用硅基流动（即使没有密钥也会报错）
|
||
if not models:
|
||
models = [
|
||
("siliconflow", "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"),
|
||
("siliconflow", "Qwen/Qwen2.5-7B-Instruct")
|
||
]
|
||
|
||
self.clients = []
|
||
for provider_str, model in models:
|
||
try:
|
||
provider = LLMProvider(provider_str)
|
||
self.clients.append(LLMClient(provider, model, self.config))
|
||
except Exception as e:
|
||
logger.warning(f"初始化{provider_str}客户端失败: {e}，跳过该模型")
|
||
|
||
self.prompt_builder = REPromptBuilder()
|
||
|
||
def extract(
|
||
self,
|
||
sentence: str,
|
||
entities: List[Dict[str, str]],
|
||
syntax_info: Dict = None,
|
||
context: str = None,
|
||
use_verification: bool = True,
|
||
) -> List[Dict[str, Any]]:
|
||
"""
|
||
从句子中提取关系（多模型并行，支持二次对话验证）
|
||
|
||
Args:
|
||
sentence: 输入句子
|
||
entities: 已识别的实体列表，格式：[{"text": "实体文本", "type": "实体类型"}, ...]
|
||
syntax_info: 句法分析结果
|
||
context: 上下文段落
|
||
use_verification: 是否使用二次对话验证
|
||
|
||
Returns:
|
||
结果列表，每个包含模型名称和三元组列表
|
||
"""
|
||
# 格式化实体列表（用于Prompt）
|
||
formatted_entities = self._format_entities_for_prompt(entities)
|
||
|
||
# 格式化句法信息
|
||
formatted_syntax = self._format_syntax_info(syntax_info) if syntax_info else "无句法信息"
|
||
|
||
# 构建Prompt模板
|
||
template = self.prompt_builder.build_template(
|
||
entities=entities,
|
||
syntax_info=formatted_syntax,
|
||
context=context or ""
|
||
)
|
||
|
||
# 构建第一阶段消息
|
||
system_msg = template.task_description["system"]
|
||
user_msg = template.task_description["user"]
|
||
|
||
# 构建完整的用户消息
|
||
input_data = f"""已识别实体：
|
||
{formatted_entities}
|
||
|
||
句子文本：
|
||
{sentence}
|
||
|
||
句法分析结果：
|
||
{formatted_syntax}
|
||
|
||
上下文段落：
|
||
{context or "无上下文"}"""
|
||
|
||
full_user_msg = user_msg
|
||
if template.candidate_targets:
|
||
full_user_msg += "\n\n" + template.build_candidate_targets_section()
|
||
if template.task_examples:
|
||
full_user_msg += "\n\n" + template.build_task_examples_section()
|
||
if template.task_emphasis:
|
||
full_user_msg += "\n\n" + template.build_task_emphasis_section()
|
||
full_user_msg += f"\n\n输入数据：\n{input_data}"
|
||
|
||
messages = [
|
||
{"role": "system", "content": system_msg},
|
||
{"role": "user", "content": full_user_msg}
|
||
]
|
||
|
||
# 并行调用多个模型
|
||
results = []
|
||
for i, client in enumerate(self.clients):
|
||
try:
|
||
model_name = f"model_{i+1}"
|
||
logger.info(f"使用模型 {model_name} 进行关系抽取...")
|
||
|
||
# 第一阶段：初始提取
|
||
first_result = client.chat(messages)
|
||
|
||
# 解析第一阶段结果
|
||
first_triplets = self._parse_triplets(first_result)
|
||
|
||
# 第二阶段：验证和优化
|
||
if use_verification and first_triplets:
|
||
verification_prompt = self.prompt_builder.build_verification_prompt()
|
||
|
||
# 移除assistant消息，将第一阶段结果整合到user消息中（符合硅基流动API要求）
|
||
second_messages = [
|
||
{"role": "system", "content": "你需要验证和优化之前提取的三元组结果。请特别关注语义逻辑的正确性、遗漏关系的补全，以及利用上下文补全省略的成分。"},
|
||
{"role": "user", "content": f"{verification_prompt}\n\n原始输入：\n{input_data}\n\n第一阶段结果：\n{first_result}\n\n请验证并优化这个结果，确保输出为有效的JSON格式。"}
|
||
]
|
||
|
||
second_result = client.chat(second_messages)
|
||
final_triplets = self._parse_triplets(second_result)
|
||
else:
|
||
final_triplets = first_triplets
|
||
|
||
results.append({
|
||
"model": model_name,
|
||
"triplets": final_triplets,
|
||
"first_triplets": first_triplets,
|
||
})
|
||
|
||
except Exception as e:
|
||
logger.error(f"模型 {i+1} 调用失败: {e}")
|
||
results.append({
|
||
"model": f"model_{i+1}",
|
||
"triplets": [],
|
||
"error": str(e)
|
||
})
|
||
|
||
return results
|
||
|
||
def _format_entities_for_prompt(self, entities: List[Dict[str, str]]) -> str:
|
||
"""格式化实体列表用于Prompt"""
|
||
if not entities:
|
||
return "无识别实体"
|
||
|
||
lines = []
|
||
for entity in entities:
|
||
entity_type = entity.get("type", "未知类型")
|
||
entity_text = entity.get("text", "")
|
||
lines.append(f"- {entity_type}: {entity_text}")
|
||
|
||
return "\n".join(lines)
|
||
|
||
def _format_syntax_info(self, syntax_info: Dict) -> str:
|
||
"""格式化句法信息"""
|
||
if not syntax_info:
|
||
return "无句法信息"
|
||
|
||
formatted = []
|
||
if "dependency" in syntax_info:
|
||
deps = syntax_info["dependency"]
|
||
formatted.append("依存关系：")
|
||
for dep in deps[:5]: # 只显示前5个
|
||
formatted.append(f" {dep.get('word', '')} <-{dep.get('deprel', '')}- {dep.get('head', '')}")
|
||
|
||
return "\n".join(formatted)
|
||
|
||
def _parse_triplets(self, result: str) -> List[Dict[str, str]]:
|
||
"""解析三元组结果"""
|
||
# 尝试解析JSON
|
||
try:
|
||
triplets = json.loads(result)
|
||
if isinstance(triplets, list):
|
||
# 验证三元组格式
|
||
valid_triplets = []
|
||
for t in triplets:
|
||
if isinstance(t, dict) and "head" in t and "relation" in t and "tail" in t:
|
||
valid_triplets.append({
|
||
"head": str(t.get("head", "")),
|
||
"relation": str(t.get("relation", "")),
|
||
"tail": str(t.get("tail", ""))
|
||
})
|
||
return valid_triplets
|
||
elif isinstance(triplets, dict):
|
||
# 可能是嵌套结构，尝试提取
|
||
return []
|
||
return []
|
||
except json.JSONDecodeError:
|
||
# 尝试从文本中提取JSON
|
||
import re
|
||
json_match = re.search(r'\[.*\]', result, re.DOTALL)
|
||
if json_match:
|
||
try:
|
||
triplets = json.loads(json_match.group())
|
||
if isinstance(triplets, list):
|
||
return triplets
|
||
except:
|
||
pass
|
||
logger.warning(f"无法解析三元组结果: {result[:100]}")
|
||
return []
|
||
|
||
|