import os import psycopg2 from psycopg2 import pool from typing import List, Dict, Optional import re from llm_client import LLMClient class TextSliceProcessor: def __init__(self, text_dir: str, prompt_template: str, db_host="localhost", db_port=5432, db_name="postgres", db_user="postgres", db_password="your_password", min_conn=1, max_conn=10, api_key: str = None, model: str = None, base_url: str = None): """ 初始化文本切片处理器 Args: text_dir: 文本目录路径 prompt_template: 提示词模板 db_host: PostgreSQL数据库主机 db_port: PostgreSQL数据库端口 db_name: PostgreSQL数据库名称 db_user: PostgreSQL用户名 db_password: PostgreSQL密码 min_conn: 连接池最小连接数 max_conn: 连接池最大连接数 """ self.text_dir = text_dir self.prompt_template = prompt_template # 初始化LLM客户端 self.llm_client = LLMClient(api_key=api_key, model=model, base_url=base_url) # 创建数据库连接池 self.pool = psycopg2.pool.SimpleConnectionPool( min_conn, max_conn, host=db_host, port=db_port, database=db_name, user=db_user, password=db_password ) # 获取一个连接和游标 self.conn = self.pool.getconn() self.cursor = self.conn.cursor() # 用于缓存实体信息的字典 self.entity_cache = {} # 初始化错误日志列表 self.error_logs = [] def __del__(self): """ 析构函数，确保关闭数据库连接和连接池 """ if hasattr(self, 'cursor') and self.cursor: self.cursor.close() if hasattr(self, 'conn') and self.conn: self.pool.putconn(self.conn) if hasattr(self, 'pool') and self.pool: self.pool.closeall() def _get_sorted_files(self, directory: str) -> List[str]: """ 获取目录下所有txt文件的列表，按章节和节数字顺序返回 Args: directory: 目录路径 Returns: List[str]: 文件路径列表 """ # 存储目录路径和对应的章节数字 chapters_with_numbers = [] # 首先获取所有章节目录并排序 for root, dirs, _ in os.walk(directory): for dir_name in dirs: # 提取章节数字 match = re.search(r'第(\d+)章', dir_name) if match: chapter_num = int(match.group(1)) chapter_path = os.path.join(root, dir_name) chapters_with_numbers.append((chapter_num, chapter_path)) # 按章节数字排序 chapters_with_numbers.sort(key=lambda x: x[0]) # 存储所有文件路径 all_files = [] # 遍历排序后的章节目录 for chapter_num, chapter_path in chapters_with_numbers: # 存储当前章节下的所有节 sections_with_numbers = [] # 获取当前章节目录下的所有文件 for root, _, filenames in os.walk(chapter_path): for filename in filenames: if filename.endswith('.txt') and '_split_' in filename: file_path = os.path.join(root, filename) # 提取节数字 section_match = re.search(r'第(\d+)节', filename) if section_match: section_num = int(section_match.group(1)) sections_with_numbers.append((chapter_num, section_num, file_path)) else: # 如果没有找到节数字，将其放在当前章节的最后 sections_with_numbers.append((chapter_num, float('inf'), file_path)) # 先按章节号排序，再按节号排序 sections_with_numbers.sort(key=lambda x: (x[0], x[1])) all_files.extend([file_path for _, _, file_path in sections_with_numbers]) return all_files def _read_text_file(self, file_path: str) -> str: """ 读取文本文件内容 Args: file_path: 文件路径 Returns: str: 文件内容 """ try: with open(file_path, 'r', encoding='utf-8') as f: lines = f.readlines() if len(lines) > 1: # 跳过第一行，返回其余内容 return ''.join(lines[1:]).strip() return '' except Exception as e: error_msg = f"读取文件 {file_path} 失败: {str(e)}" print(error_msg) self.error_logs.append({"error": error_msg, "file": file_path}) return "" def _create_or_get_node(self, node_name: str, category: str) -> int: """ 创建或获取节点 Args: node_name: 节点名称 category: 节点类别 Returns: int: 节点ID """ # 检查缓存 cache_key = f"{node_name}_{category}" if cache_key in self.entity_cache: return self.entity_cache[cache_key] # 查询是否存在 self.cursor.execute( "SELECT id FROM kg_nodes WHERE name = %s AND category = %s", (node_name, category) ) result = self.cursor.fetchone() if result: node_id = result[0] else: # 创建新节点 self.cursor.execute( "INSERT INTO kg_nodes (name, category, status, version) " "VALUES (%s, %s, 0, '1.3') RETURNING id", (node_name, category) ) node_id = self.cursor.fetchone()[0] # 更新缓存 self.entity_cache[cache_key] = node_id return node_id def _create_relationship(self, text_content: str, prev_text: str, file_path: str) -> None: """ 创建文本关系并存储到数据库 Args: text_content: 当前文本内容 prev_text: 前置文本内容 file_path: 文件路径 """ try: # 构建完整的提示词内容 full_content = f"{self.prompt_template}用户输入的\"前置医学文本\"为：\n{prev_text if prev_text else '无'}\n用户输入的\"任务医学文本\"为：\n{text_content}" # 调用LLM获取结果 response = self.llm_client.chat_completion([{"role": "user", "content": full_content}]) if not hasattr(response, 'choices') or len(response.choices) == 0: raise Exception("未收到有效回复") # 获取LLM返回的结果 llm_result = response.choices[0].message.content # 开始事务 self.cursor.execute("BEGIN") # 处理每一行结果 for line in llm_result.strip().split('\n'): if not line.strip(): continue # 解析结果 parts = line.strip().split('---') if len(parts) != 5: continue start_node_name, relation_name, end_node_name, start_category, end_category = parts # 创建或获取起始节点和结束节点 start_node_id = self._create_or_get_node(start_node_name, start_category) end_node_id = self._create_or_get_node(end_node_name, end_category) # 创建关系 self.cursor.execute( "INSERT INTO kg_edges (category, src_id, dest_id, name, status, version) " "VALUES (%s, %s, %s, %s, 0, '1.3') RETURNING id", (relation_name, start_node_id, end_node_id, relation_name) ) edge_id = self.cursor.fetchone()[0] # 添加文件路径属性 self.cursor.execute( "INSERT INTO kg_props (category, ref_id, prop_name, prop_value, type, prop_title) " "VALUES (4, %s, 'file_path', %s, 1, 'file_path')", (edge_id, file_path) ) # 提交事务 self.conn.commit() print(f"处理文件完成: {file_path}") except Exception as e: self.conn.rollback() error_msg = f"处理文件 {file_path} 时出错: {str(e)}" print(error_msg) self.error_logs.append({"error": error_msg, "file": file_path}) def process_text_files(self): """ 处理目录下的所有文本文件 """ try: # 记录处理情况 self.processed_files = 0 self.total_files = 0 # 递归处理目录 self._process_directory(self.text_dir) # 输出处理总结 print("\n处理完成！") print(f"总文件数: {self.total_files}") print(f"成功处理: {self.processed_files}") print(f"处理失败: {len(self.error_logs)}") if self.error_logs: print("\n失败文件列表:") for log in self.error_logs: print(f"文件: {log['file']}") print(f"错误: {log['error']}\n") except Exception as e: print(f"批量处理失败: {str(e)}") def _process_directory(self, directory: str, prev_text: str = "") -> None: """ 递归处理目录下的文件 Args: directory: 目录路径 prev_text: 前置文本内容 """ try: # 获取目录下的所有文件和子目录 items = sorted(os.listdir(directory)) # 收集当前目录下的所有txt文件 txt_files = [] for item in items: item_path = os.path.join(directory, item) if os.path.isfile(item_path) and item.endswith('.txt') and '_split_' in item: # 提取split序号 match = re.search(r'split_(\d+)', item) if match: split_num = int(match.group(1)) txt_files.append((split_num, item_path)) # 按split序号排序文件 txt_files.sort(key=lambda x: x[0]) # 处理排序后的文件 for split_num, item_path in txt_files: self.total_files += 1 print(f"\n开始处理文件 [{self.processed_files + 1}/{self.total_files}]: {os.path.basename(item_path)}") try: # 读取当前文件内容 text_content = self._read_text_file(item_path) if not text_content: continue # 如果是第一个文件（split_0），前置文本为'无' if split_num == 0: current_prev_text = '无' else: # 获取前一个文件的路径 prev_path = item_path.replace(f'split_{split_num}', f'split_{split_num-1}') # 如果前一个文件存在，读取其内容作为前置文本 if os.path.exists(prev_path): current_prev_text = self._read_text_file(prev_path) else: current_prev_text = prev_text # 创建关系 self._create_relationship(text_content, current_prev_text, item_path) self.processed_files += 1 # 更新前置文本 prev_text = text_content except Exception as e: error_msg = f"处理文件 {item_path} 时发生错误: {str(e)}" print(error_msg) self.error_logs.append({"error": error_msg, "file": item_path}) continue # 然后递归处理子目录 for item in items: item_path = os.path.join(directory, item) if os.path.isdir(item_path): # 检查目录名是否包含章节信息 if re.search(r'第(\d+)章', item): self._process_directory(item_path, prev_text) except Exception as e: error_msg = f"处理目录 {directory} 时发生错误: {str(e)}" print(error_msg) self.error_logs.append({"error": error_msg, "file": directory}) def main(): # 解析命令行参数 import argparse parser = argparse.ArgumentParser(description='文本切片处理工具') parser.add_argument('--text_dir', type=str, required=False, default='/Users/ycw/work/qwen32b知识抽取/《急诊医学（第2版）》', help='文本目录路径') prompt = """ ##角色任务 你是一个医学专家，且精通中文文本的实体和关系标注（自然语言处理），用户将提供两个医学文本：“前置医学文本”和“任务医学文本”，你的任务是从用户输入的“任务医学文本”中抽取所有的“关系型”知识要点，构建符合「头部实体---关系---尾部实体」的三元组集合。 ##请严格依次执行以下步骤及其细节要求： 第一步，通读“前置医学文本”和“任务医学文本”，对“任务医学文本”形成一个整体的理解，搞明白其围绕的主题是什么？主要讲的是什么？主要说了哪几块或几点内容？。此处“前置医学文本”是“任务医学文本”的上文语义背景，其有助于你对“任务医学文本”语义的理解，且在你提取三元组时，其可能作为参考或辅助信息。 第二步【实体识别】： 规则1：可参考但不限于如下医学实体类型：症状、体格检查项目（简称：体征项）、体格检查项目结果（简称：“体征”或“体征结果”）、检查、检查结果、疾病、药品、手术操作、其他治疗、科室、人群、人体结构或部位、医疗器械、食物、病理机制等。 规则2：复合实体要拆分：如“老年人和免疫功能低下者”不能作为一个实体，需拆分为："老年人"、"免疫功能低下者"两个实体。 规则3：文字中省略了，但语义中暗含的文字内容要进行语义补全：如文本：“胸、腹部检查”中，这里的“胸”其实是指“胸部检查”，按语义补全后应该有两个实体：胸部检查、腹部检查。同理，文本：“注意有无心内膜炎、心肌炎、心包炎体征”中，应该有三个实体：内膜炎体征、心肌炎体征、心包炎体征；文本：“有无肝脏和脾脏肿大”中，应该有两个实体：肝脏肿大、脾脏肿大。 第三步【关系构建】： 规则4：可参考的“基础关系类型”如下（但不仅仅限于这些）：属于（是）、包括（包含）、导致（的结果）、是由…导致（的结果）、的原因是、是…的原因、的病因是、是…的病因、基于（基础是）、是…的基础、推荐、被推荐于、区别于、相似于、关联、疾病常关联、疾病可推荐等等。 规则5：根据上下文整体语义分析并确定出两个实体之间的“基础关系类型”后，可能还不足以描述清楚两个实体间的“详细关系”或“精准的关系”，所以你要尽量构建字数更多的详细的“关系”，需要依据原文中真实的语义，在“基础关系”中增加“修饰词”、“限定条件词”等语义描述，避免关系构建太过粗糙，避免原文关系语义信息的衰减或丢失。 规则6：语义中属于该三元组关系的修饰词、限定条件词等，需要融入到该“关系”中（如：可能、少数、30%、多数、显著、轻微、手术后、治疗无效时、满足XX条件时等），如果归属于该“关系”，则需要组合到该“关系”中使得关系更丰满（如：“可能导致”、“少数由…导致”、“手术后导致”、“满足XX条件时，推荐”等等）。 规则7：“尾部实体”不能为复合实体或并排结构，必须合理的拆分为多个三元组。 规则8：强制质量检测：“头部实体---关系类型---尾部实体”中，三者必须能组合成语法通顺的一句话，不能有语病，且关系指向不能模糊或错误。如三元组1：“A---的原因是---B”中，组合成一句话“A的原因是B”没有语病，且关系明确：B是原因，A是B原因导致的结果；如三元组2：“A---是…的原因---B”中，组合成一句话“A是B的原因”没有语病，且关系明确：A是原因，B是A原因导致的结果。注意：不是所有的关系都是镜像对称的，很多关系是单向关系，请根据语义关系的方向选择正确的关系类型，否则语义会完全相反。 ##输出格式要求： 对于每个三元组关系，提取以下信息：头部实体的名称（source_entity）、尾部实体的名称（target_entity）、关系（relationship_type）、头部实体的类型（source_entity_type）、尾部实体的类型（target_entity_type），最后这些信息整理成如下格式的一个“字符串”：“头部实体的名称---关系---尾部实体的名称---头部实体的类型---尾部实体的类型”，不同的三元组“字符串”之间用换行符号连接起来。 ##质量红线 1.禁止出现四元组或嵌套结构 2.禁止合并多个“差异点”或“并列实体”到单个三元组 3.每个三元组必须保留原文核心逻辑 ##示例： “前置医学文本”为： 第一节，自发性气胸 “任务医学文本”为： "三、诊断要点 自发性气胸通过胸部 X 线片确立诊断，条件允许时，应选择直立位拍片。 1. 既往胸部 X 线检查无明显病变或有 COPD、肺结核、哮喘等肺部基础病变。 2. 突发一侧胸痛伴不同程度的胸闷、呼吸困难。患侧胸廓饱满、呼吸运动减弱，叩诊呈鼓音，肝、 肺浊音界消失，听诊呼吸音减弱，甚至消失。 3. 发病时胸部 X 线影像学检查是诊断气胸最为准确和可靠的方法。 典型自发性气胸诊断不难。继发性气胸病人可因原有基础疾病而影响诊断，因此，对临床不能用 其他原因解释或经急诊处理呼吸困难无改善者，需考虑自发性气胸的可能。因病情危重不能立即行 胸部 X 线检查时，可在胸腔积气体征最明显处进行诊断性穿刺。" 输出为： "自发性气胸---通过...确立诊断---胸部X线片---疾病---检查 自发性气胸---条件允许时推荐选择---直立位拍片---疾病---检查操作 自发性气胸---诊断要点需结合---既往胸部X线检查无明显病变---疾病---检查结果 自发性气胸---诊断要点需结合---或有肺部基础病变---疾病---检查结果 肺部基础病变---的原因包括---COPD---检查结果---疾病 肺部基础病变---的原因包括---肺结核---检查结果---疾病 肺部基础病变---的原因包括---哮喘---检查结果---疾病 自发性气胸---诊断要点需结合---突发一侧胸痛伴不同程度的胸闷 自发性气胸---诊断要点需结合---突发一侧胸痛伴不同程度的呼吸困难 突发一侧胸痛---伴发---胸闷---症状---症状 突发一侧胸痛---伴发---呼吸困难---症状---症状 自发性气胸---患侧表现---患侧胸廓饱满---疾病---体征结果 自发性气胸---患侧表现---呼吸运动减弱---疾病---体征结果 自发性气胸---患侧表现---叩诊呈鼓音---疾病---体征结果 自发性气胸---体征结果---肝脏浊音界消失---疾病---体征结果 自发性气胸---体征结果---肺浊音界消失---疾病---体征结果 自发性气胸---患侧表现---肺浊音界消失---疾病---体征结果 自发性气胸---患侧表现---听诊呼吸音减弱---疾病---体征结果 自发性气胸---患侧表现---听诊呼吸音消失---疾病---体征结果 胸部X线影像学检查---是诊断...最准确可靠方法---气胸---检查---疾病 继发性气胸病人---可因...被影响诊断---原有基础疾病---人群---疾病 临床不能用其他原因解释的呼吸困难者---需考虑可能为---自发性气胸---人群---疾病 急诊处理呼吸困难无改善者---需考虑可能为---自发性气胸---人群---疾病 病情危重不能立即行胸部X线检查者---推荐进行---诊断性穿刺---人群---检查操作 诊断性穿刺---实施部位---胸腔积气体征最明显处---检查操作---人体部位" """ parser.add_argument('--prompt', type=str, required=False, default=prompt, help='提示词模板') args = parser.parse_args() # 设置PostgreSQL连接参数 db_host = "173.18.12.203" db_port = 5432 db_name = "kg_ycw" db_user = "knowledge" db_password = "qwer1234." # 请替换为实际的密码 try: # 创建处理器实例 processor = TextSliceProcessor( text_dir=args.text_dir, prompt_template=args.prompt, db_host=db_host, db_port=db_port, db_name=db_name, db_user=db_user, db_password=db_password ) # 执行批量处理 processor.process_text_files() except Exception as e: print(f"错误: {str(e)}") if __name__ == "__main__": main()