""" 关系抽取模块 使用多个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 []