# -*- coding: utf-8 -*- """ 将重新抽取的结果合并到现有的节点和关系文件中 """ import json import pandas as pd from pathlib import Path import sys # 添加模块路径 sys.path.append(str(Path(__file__).parent / 'src')) from data_processing.entity_normalizer import EntityNormalizer from data_processing.relationship_builder import RelationshipBuilder def merge_retry_results(): """合并重新抽取的结果""" # 文件路径 retry_file = Path(__file__).parent / 'output' / 'retry_projects.json' nodes_file = Path(__file__).parent / 'output' / 'nodes_llm.csv' rels_file = Path(__file__).parent / 'output' / 'rels_llm.csv' ontology_file = Path(__file__).parent / 'config' / 'entity_ontology.yaml' print("=== 加载数据 ===") # 读取重新抽取的结果 with open(retry_file, 'r', encoding='utf-8') as f: retry_data = json.load(f) print(f"重新抽取的项目数: {len(retry_data)}") # 读取现有的节点和关系 existing_nodes_df = pd.read_csv(nodes_file, encoding='utf-8-sig') existing_rels_df = pd.read_csv(rels_file, encoding='utf-8-sig') print(f"现有节点数: {len(existing_nodes_df)}") print(f"现有关系数: {len(existing_rels_df)}") # 初始化规范化器和关系构建器 normalizer = EntityNormalizer(str(ontology_file)) builder = RelationshipBuilder(str(ontology_file)) # 规范化重新抽取的实体 print("\n=== 规范化实体 ===") # 收集所有需要规范化的实体 all_entities_to_normalize = [] for item in retry_data: result = item['extraction_result'] entities = result.get('entities', []) # 添加项目ID以便跟踪 for entity in entities: entity['_project_id'] = item['project_id'] all_entities_to_normalize.extend(entities) # 批量规范化 extraction_results = [] for item in retry_data: extraction_results.append(item['extraction_result']) normalizer.normalize_batch(extraction_results) # 生成新的实体节点 new_entity_nodes = normalizer.get_entity_nodes() print(f"新生成实体节点数: {len(new_entity_nodes)}") # 添加项目节点 new_project_nodes = [] for item in retry_data: new_project_nodes.append({ 'id': item['project_id'], 'label': item['project_name'], 'type': 'ICH_Project', 'properties': '{}' }) print(f"新增项目节点数: {len(new_project_nodes)}") # 构建关系 print("\n=== 构建关系 ===") all_new_relationships = [] # 获取当前所有节点的ID到标签映射（用于查找目标实体ID） all_nodes_for_lookup = pd.concat([ existing_nodes_df, pd.DataFrame(new_entity_nodes) ], ignore_index=True) # 创建名称->ID的映射字典 name_to_id = {} for idx, row in all_nodes_for_lookup.iterrows(): if row['type'] != 'ICH_Project': # 只映射实体节点 name_to_id[(row['type'], row['label'])] = row['id'] for item in retry_data: project_id = item['project_id'] result = item['extraction_result'] entities = result.get('entities', []) relationships = result.get('relationships', []) # 构建实体名称到ID的映射（仅限当前项目） entity_name_to_id = {} for entity in entities: # 实体名称可能在name字段或attributes.name字段 entity_name = entity.get('name') or entity.get('attributes', {}).get('name', '') entity_type = entity.get('type', '') normalized_name = normalizer.normalize_text(entity_name) # 在规范化器的映射中查找（键是normalized_text，不是元组） entity_id = normalizer.text_to_id_map.get(normalized_name) if not entity_id: # 在所有节点中查找（可能是已存在的实体） entity_id = name_to_id.get((entity_type, entity_name)) if entity_id: entity_name_to_id[entity_name] = entity_id # 手动构建关系 for rel in relationships: # 关系中的字段名是source_entity和target_entity source_name = rel.get('source_entity') or rel.get('source') target_name = rel.get('target_entity') or rel.get('target') rel_type = rel.get('type') rel_props = rel.get('properties', {}) # 查找源实体ID if source_name == project_id: source_id = project_id else: # 先尝试在当前项目实体中查找 source_id = entity_name_to_id.get(source_name) if not source_id: # 在规范化器的映射中查找 normalized_name = normalizer.normalize_text(source_name) source_id = normalizer.text_to_id_map.get(normalized_name) # 查找目标实体ID # 先尝试在当前项目实体中查找 target_id = entity_name_to_id.get(target_name) if not target_id: # 在规范化器的映射中查找 normalized_name = normalizer.normalize_text(target_name) target_id = normalizer.text_to_id_map.get(normalized_name) # 如果都找到了，添加关系 if source_id and target_id: all_new_relationships.append({ 'source': source_id, 'target': target_id, 'type': rel_type, 'properties': json.dumps(rel_props, ensure_ascii=False) if rel_props else '{}' }) print(f"项目 {project_id}: {len([r for r in all_new_relationships if r['source'] == project_id])} 条关系") print(f"新增关系总数: {len(all_new_relationships)}") # 合并节点 print("\n=== 合并节点 ===") # 过滤掉已经存在的项目节点 existing_project_ids = set(existing_nodes_df[existing_nodes_df['type'] == 'ICH_Project']['id'].tolist()) new_project_nodes_filtered = [n for n in new_project_nodes if n['id'] not in existing_project_ids] all_nodes = pd.concat([ existing_nodes_df, pd.DataFrame(new_entity_nodes), pd.DataFrame(new_project_nodes_filtered) ], ignore_index=True) print(f"合并后节点数: {len(all_nodes)} (新增 {len(new_entity_nodes) + len(new_project_nodes_filtered)} 个)") # 合并关系 print("\n=== 合并关系 ===") all_rels = pd.concat([ existing_rels_df, pd.DataFrame(all_new_relationships) ], ignore_index=True) print(f"合并后关系数: {len(all_rels)} (新增 {len(all_new_relationships)} 条)") # 保存结果 print("\n=== 保存结果 ===") all_nodes.to_csv(nodes_file, index=False, encoding='utf-8-sig') all_rels.to_csv(rels_file, index=False, encoding='utf-8-sig') print(f"节点已保存: {nodes_file}") print(f"关系已保存: {rels_file}") # 验证 print("\n=== 验证 ===") for item in retry_data: project_id = item['project_id'] rel_count = len(all_rels[all_rels['source'] == project_id]) print(f"{project_id}: {rel_count} 条关系") print("\n完成！") if __name__ == '__main__': merge_retry_results()