591 lines
20 KiB
Python
591 lines
20 KiB
Python
import json
|
||
import fitz # PyMuPDF
|
||
import os
|
||
import fitz
|
||
from collections import Counter
|
||
import uuid
|
||
from kg_build.Prompt import Prompt_compose_tujie1,Prompt_compose_tujie2
|
||
from kg_build.extract_util import safe_json_loads,highlight_pdf_file,get_content_bbox
|
||
# from modelsAPI.model_api import OpenaiAPI
|
||
import re
|
||
from openai import AsyncOpenAI
|
||
import asyncio
|
||
import json
|
||
import re
|
||
import copy
|
||
import pandas as pd
|
||
from io import StringIO
|
||
import json
|
||
from collections import defaultdict
|
||
import traceback
|
||
|
||
|
||
def is_tujie_compose_table_by_first_tr(table_body: str) -> bool:
|
||
if not isinstance(table_body, str) or not table_body.strip():
|
||
return False
|
||
|
||
text_without_tags = re.sub(r'<[^>]+>', '', table_body)
|
||
text_lower = text_without_tags.lower()
|
||
|
||
keyword_groups = [
|
||
["组成编码", "名称", "是否为关重件", "是否为寿命件"],
|
||
["组成编码", "名称", "型号/零件号", "重量(kg)", "数量"]
|
||
]
|
||
|
||
for group in keyword_groups:
|
||
if all(kw in text_lower for kw in group):
|
||
return True
|
||
|
||
return False
|
||
|
||
|
||
def is_tujie_compose_mulus_table_by_first_tr(table_body: str) -> bool:
|
||
if not isinstance(table_body, str) or not table_body.strip():
|
||
return False
|
||
|
||
text_without_tags = re.sub(r'<[^>]+>', '', table_body)
|
||
text_lower = text_without_tags.lower()
|
||
|
||
keyword_groups = [
|
||
["组成编码", "名称", "型号/零件号", "重量(kg)", "数量"]
|
||
]
|
||
|
||
for group in keyword_groups:
|
||
if all(kw in text_lower for kw in group):
|
||
return True
|
||
|
||
return False
|
||
|
||
|
||
def get_tujiecompose_chunk_bbox_and_removed_data(data):
|
||
"""
|
||
返回：
|
||
- 符合条件的图解组成表分组（每组 [前前, 前, 当前]）
|
||
- 从原始 data 中移除了所有相关记录后的新 data 列表
|
||
"""
|
||
title = "组成"
|
||
record_map = {item["id"]: item for item in data}
|
||
|
||
results = []
|
||
ids_to_remove = set()
|
||
|
||
for ins in data:
|
||
if ins.get("type") == "table":
|
||
caption_has_title = any(title in caption for caption in ins.get("table_caption", []))
|
||
|
||
table_body = ins.get("table_body", "")
|
||
body_matches_keywords = is_tujie_compose_table_by_first_tr(table_body)
|
||
|
||
if caption_has_title and body_matches_keywords:
|
||
current_id = ins["id"]
|
||
prev_ids = [current_id - 2, current_id - 1]
|
||
|
||
group_ids = [pid for pid in prev_ids if pid >= 1] + [current_id]
|
||
ids_to_remove.update(group_ids)
|
||
|
||
group = [record_map.get(pid) for pid in prev_ids if pid >= 1] + [ins]
|
||
results.append(group)
|
||
|
||
filtered_data = [item for item in data if item["id"] not in ids_to_remove]
|
||
return results, filtered_data
|
||
|
||
|
||
def get_tujiecompose_mulu_chunk_bbox_and_removed_data(data):
|
||
title = "组成"
|
||
record_map = {item["id"]: item for item in data}
|
||
|
||
results = []
|
||
ids_to_remove = set()
|
||
|
||
for ins in data:
|
||
if ins.get("type") == "table":
|
||
caption_has_title = any(title in caption for caption in ins.get("table_caption", []))
|
||
|
||
table_body = ins.get("table_body", "")
|
||
body_matches_keywords = is_tujie_compose_mulus_table_by_first_tr(table_body)
|
||
|
||
if body_matches_keywords:
|
||
current_id = ins["id"]
|
||
prev_ids = [current_id - 2, current_id - 1]
|
||
|
||
group_ids = [pid for pid in prev_ids if pid >= 1] + [current_id]
|
||
ids_to_remove.update(group_ids)
|
||
|
||
group = [record_map.get(pid) for pid in prev_ids if pid >= 1] + [ins]
|
||
results.append(group)
|
||
|
||
filtered_data = [item for item in data if item["id"] not in ids_to_remove]
|
||
return results, filtered_data
|
||
|
||
|
||
def format_compose_code(val):
|
||
"""将 int/float 型组成编码转换为偶数位字符串，不足偶数位则前补 0"""
|
||
if pd.isna(val):
|
||
return None
|
||
if isinstance(val, (int, float)):
|
||
s = str(int(val))
|
||
if len(s) % 2 != 0:
|
||
s = '0' + s
|
||
return s
|
||
return str(val)
|
||
|
||
|
||
def merge_compose(records):
|
||
"""
|
||
合并相同组成编码的记录。
|
||
特殊处理：将 knowledge_source 字段的多个值解析并合并到一个新的扁平列表中。
|
||
"""
|
||
grouped_data = defaultdict(list)
|
||
|
||
for record in records:
|
||
code = record.get('组成编码')
|
||
if code:
|
||
grouped_data[code].append(record)
|
||
|
||
merged_data = []
|
||
|
||
for key, group in grouped_data.items():
|
||
merged_record = group[0].copy()
|
||
|
||
all_sources = []
|
||
|
||
for record in group:
|
||
source_val = record.get('knowledge_source')
|
||
if source_val:
|
||
try:
|
||
if isinstance(source_val, str):
|
||
parsed_source = json.loads(source_val)
|
||
else:
|
||
parsed_source = source_val
|
||
|
||
if isinstance(parsed_source, list):
|
||
all_sources.extend(parsed_source)
|
||
else:
|
||
all_sources.append(parsed_source)
|
||
except (json.JSONDecodeError, TypeError):
|
||
all_sources.append(source_val)
|
||
|
||
if all_sources:
|
||
merged_record['knowledge_source'] = json.dumps(all_sources, ensure_ascii=False)
|
||
else:
|
||
merged_record['knowledge_source'] = "[]"
|
||
|
||
for record in group[1:]:
|
||
for k, v in record.items():
|
||
if k == 'knowledge_source':
|
||
continue
|
||
|
||
existing = merged_record.get(k)
|
||
|
||
if v and existing and existing != v:
|
||
if isinstance(existing, list):
|
||
if v not in existing:
|
||
existing.append(v)
|
||
else:
|
||
merged_record[k] = [existing, v]
|
||
elif not existing and v:
|
||
merged_record[k] = v
|
||
|
||
merged_data.append(merged_record)
|
||
|
||
return merged_data
|
||
|
||
|
||
def remove_fields(data_list, fields_to_remove):
|
||
"""
|
||
遍历列表中的每个字典，移除指定的字段（键）。
|
||
如果字段不存在，则跳过该字段。
|
||
"""
|
||
cleaned_data = []
|
||
|
||
for item in data_list:
|
||
new_item = item.copy()
|
||
|
||
for field in fields_to_remove:
|
||
if field in new_item:
|
||
del new_item[field]
|
||
|
||
cleaned_data.append(new_item)
|
||
|
||
return cleaned_data
|
||
|
||
|
||
# def generate_entities_and_relations(data_list):
|
||
# """
|
||
# 所有实体统一 type="设备"。
|
||
# 关系构建：若父节点在数据中不存在，自动补充占位实体，确保所有关系都成功建立。
|
||
# """
|
||
# entities = []
|
||
# relationships = []
|
||
|
||
# code_name_map = {} # {code: name}
|
||
# clean_nodes = []
|
||
|
||
# # --- 第一步：数据清洗与标准化 ---
|
||
# for item in data_list:
|
||
# code = item.get('组成编码', '')
|
||
|
||
# if not code or not isinstance(code, str):
|
||
# continue
|
||
# if not code.isdigit():
|
||
# continue
|
||
# if len(code) < 2 or len(code) % 2 != 0:
|
||
# continue
|
||
|
||
# cleaned_item = {}
|
||
# for key, value in item.items():
|
||
# if isinstance(value, list):
|
||
# cleaned_item[key] = str(value[0]) if value else ""
|
||
# elif value is None:
|
||
# cleaned_item[key] = ""
|
||
# else:
|
||
# cleaned_item[key] = value
|
||
|
||
# code_name_map[code] = cleaned_item.get('名称', f'未知_{code}')
|
||
# clean_nodes.append(cleaned_item)
|
||
|
||
# # --- 第二步：找出所有缺失的父节点，自动递归向上补全 ---
|
||
# placeholder_codes = {} # {code: name} 仅补充的占位节点
|
||
|
||
# for node in clean_nodes:
|
||
# code = node['组成编码']
|
||
# length = len(code)
|
||
|
||
# if length < 4:
|
||
# continue # 2位根节点无父节点，跳过
|
||
|
||
# current = code[:-2]
|
||
# while current:
|
||
# if current in code_name_map or current in placeholder_codes:
|
||
# break # 已存在，无需补全
|
||
# placeholder_name = f'未知_{current}'
|
||
# placeholder_codes[current] = placeholder_name
|
||
# print(f" ⚠️ 父节点缺失，自动补充占位实体: 编码={current}, 名称={placeholder_name}")
|
||
# if len(current) >= 4:
|
||
# current = current[:-2]
|
||
# else:
|
||
# break
|
||
|
||
# # 将占位节点合并进 code_name_map，供关系构建使用
|
||
# code_name_map.update(placeholder_codes)
|
||
|
||
# # --- 第三步：构建实体列表，所有节点 type="设备" ---
|
||
# # 真实节点
|
||
# for node in clean_nodes:
|
||
# entities.append({
|
||
# "type": "设备",
|
||
# "properties": node
|
||
# })
|
||
|
||
# # 占位节点（最简属性，标记 _is_placeholder 便于下游过滤）
|
||
# for code, name in placeholder_codes.items():
|
||
# entities.append({
|
||
# "type": "设备",
|
||
# "properties": {
|
||
# "组成编码": code,
|
||
# "名称": name,
|
||
# "knowledge_source": "[]",
|
||
# "_is_placeholder": True
|
||
# }
|
||
# })
|
||
|
||
# # --- 第四步：构建关系列表，保证全部成功 ---
|
||
# success_count = 0
|
||
# skip_count = 0
|
||
|
||
# for node in clean_nodes:
|
||
# code = node['组成编码']
|
||
# length = len(code)
|
||
|
||
# if length < 4 or length % 2 != 0:
|
||
# skip_count += 1
|
||
# continue # 2位根节点无父，正常跳过
|
||
|
||
# parent_code = code[:-2]
|
||
# child_name = node.get('名称', f'未知_{code}')
|
||
# parent_name = code_name_map.get(parent_code)
|
||
|
||
# if parent_name:
|
||
# relationships.append({
|
||
# "type": "包含",
|
||
# "from_entity": parent_name,
|
||
# "to_entity": child_name
|
||
# })
|
||
# success_count += 1
|
||
# else:
|
||
# # 经过第二步补全后理论上不会进入此分支
|
||
# print(f" ❌ 关系构建失败（不应出现）: 子编码={code}, 父编码={parent_code} 仍未找到")
|
||
# skip_count += 1
|
||
|
||
# print(f" 📊 关系构建完成：成功={success_count}, 跳过根节点={skip_count}")
|
||
# return entities, relationships
|
||
def generate_entities_and_relations(data_list):
|
||
"""
|
||
实体分两类：
|
||
- type="设备"：原始组成表中的真实/占位节点
|
||
- type="组成"：每个有子节点的设备自动生成的中间节点，名称为 "{设备名}的组成"
|
||
|
||
关系：
|
||
- 父节点（设备）→ 中间节点（组成）：包含
|
||
- 中间节点（组成）→ 子节点（设备）：包含
|
||
"""
|
||
entities = []
|
||
relationships = []
|
||
|
||
code_name_map = {} # {code: name}
|
||
clean_nodes = []
|
||
|
||
# ── 第一步：数据清洗与标准化 ──────────────────────────────────────
|
||
for item in data_list:
|
||
code = item.get('组成编码', '')
|
||
if not code or not isinstance(code, str):
|
||
continue
|
||
if not code.isdigit():
|
||
continue
|
||
if len(code) < 2 or len(code) % 2 != 0:
|
||
continue
|
||
|
||
cleaned_item = {}
|
||
for key, value in item.items():
|
||
if isinstance(value, list):
|
||
cleaned_item[key] = str(value[0]) if value else ""
|
||
elif value is None:
|
||
cleaned_item[key] = ""
|
||
else:
|
||
cleaned_item[key] = value
|
||
|
||
code_name_map[code] = cleaned_item.get('名称', f'未知_{code}')
|
||
clean_nodes.append(cleaned_item)
|
||
|
||
# ── 第二步：自动补全缺失的父节点占位实体 ─────────────────────────
|
||
placeholder_codes = {} # {code: name}
|
||
|
||
for node in clean_nodes:
|
||
code = node['组成编码']
|
||
current = code[:-2]
|
||
while current:
|
||
if current in code_name_map or current in placeholder_codes:
|
||
break
|
||
placeholder_name = f'未知_{current}'
|
||
placeholder_codes[current] = placeholder_name
|
||
print(f" ⚠️ 父节点缺失，自动补充占位实体: 编码={current}, 名称={placeholder_name}")
|
||
current = current[:-2] if len(current) >= 4 else ""
|
||
|
||
code_name_map.update(placeholder_codes)
|
||
|
||
# ── 第三步：构建"设备"实体（真实节点 + 占位节点）────────────────
|
||
for node in clean_nodes:
|
||
entities.append({"type": "设备", "properties": node})
|
||
|
||
for code, name in placeholder_codes.items():
|
||
entities.append({
|
||
"type": "设备",
|
||
"properties": {
|
||
"组成编码": code,
|
||
"名称": name,
|
||
"knowledge_source": "[]",
|
||
}
|
||
})
|
||
|
||
# ── 第四步：为每个"有子节点的父设备"生成"组成"中间实体 ──────────
|
||
# 收集所有实际存在父子关系的父编码
|
||
parent_codes_needed = set()
|
||
for node in clean_nodes:
|
||
code = node['组成编码']
|
||
if len(code) >= 4 and len(code) % 2 == 0:
|
||
parent_codes_needed.add(code[:-2])
|
||
|
||
# {parent_code: "组成中间实体名称"}
|
||
compose_node_map = {}
|
||
for parent_code in parent_codes_needed:
|
||
parent_name = code_name_map.get(parent_code)
|
||
if not parent_name:
|
||
print(f" ❌ 找不到父编码对应名称（不应出现）: {parent_code}")
|
||
continue
|
||
compose_name = f"{parent_name}的组成"
|
||
compose_node_map[parent_code] = compose_name
|
||
|
||
# 生成"组成"中间实体
|
||
entities.append({
|
||
"type": "组成",
|
||
"properties": {
|
||
"名称": compose_name,
|
||
"knowledge_source": "[]",
|
||
}
|
||
})
|
||
|
||
# 父设备 → 组成中间节点
|
||
relationships.append({
|
||
"type": "包含",
|
||
"from_entity": parent_name,
|
||
"to_entity": compose_name
|
||
})
|
||
|
||
# ── 第五步：构建"组成中间节点 → 子设备"关系 ──────────────────────
|
||
success_count, skip_count = 0, 0
|
||
|
||
for node in clean_nodes:
|
||
code = node['组成编码']
|
||
if len(code) < 4 or len(code) % 2 != 0:
|
||
skip_count += 1
|
||
continue # 根节点无父，正常跳过
|
||
|
||
parent_code = code[:-2]
|
||
compose_name = compose_node_map.get(parent_code)
|
||
child_name = node.get('名称', f'未知_{code}')
|
||
|
||
if compose_name:
|
||
relationships.append({
|
||
"type": "包含",
|
||
"from_entity": compose_name,
|
||
"to_entity": child_name
|
||
})
|
||
success_count += 1
|
||
else:
|
||
print(f" ❌ 找不到组成中间节点（不应出现）: 子编码={code}, 父编码={parent_code}")
|
||
skip_count += 1
|
||
|
||
print(f" 📊 关系构建完成：成功={success_count}, 跳过根节点={skip_count}")
|
||
return entities, relationships
|
||
|
||
async def get_tujiecompose_node_relation(input_pdf_path, data, prefix_url="http://192.168.0.46:59085", device="发动机", lower_entity_type="设备"):
|
||
# 初始化默认返回值（空列表）
|
||
all_entity_results = []
|
||
all_relation_results = []
|
||
devicename = device
|
||
lower_entity_type = lower_entity_type
|
||
|
||
try:
|
||
# 1. 初始化 ID（确保从 1 开始连续）
|
||
for idx, item in enumerate(data, start=1):
|
||
item["id"] = idx
|
||
|
||
# 2. 获取分组数据
|
||
results, filtered_data = get_tujiecompose_chunk_bbox_and_removed_data(data=data)
|
||
compose_mulu_result, filtered_data_mulu = get_tujiecompose_mulu_chunk_bbox_and_removed_data(data=data)
|
||
|
||
# 3. 合并所有需要处理的组
|
||
all_results = results + compose_mulu_result
|
||
|
||
# 初始化中间结果容器
|
||
all_htmltojson_results = []
|
||
|
||
print(f"共发现 {len(all_results)} 个图解组成表组，开始处理...")
|
||
|
||
for i, compose_group in enumerate(all_results):
|
||
try:
|
||
print(f"开始处理第 {i+1} 个图解组成表")
|
||
|
||
content, highlight_list = get_content_bbox(compose_group)
|
||
|
||
if not content:
|
||
print(f" ⚠️ 第 {i+1} 组未获取到有效内容，跳过")
|
||
continue
|
||
|
||
# 文本清洗
|
||
content = content.replace("香", "否")
|
||
content = content.replace("房号", "序号")
|
||
|
||
# 解析 HTML 表格
|
||
html_obj = StringIO(content)
|
||
df_list = pd.read_html(html_obj, header=0)
|
||
|
||
if not df_list:
|
||
print(f" ⚠️ 第 {i+1} 组未解析到表格，跳过")
|
||
continue
|
||
|
||
df = df_list[0]
|
||
|
||
# 清洗数据：删除全为空的行
|
||
df.dropna(how='all', inplace=True)
|
||
|
||
# 格式化组成编码
|
||
if '组成编码' in df.columns and pd.api.types.is_numeric_dtype(df['组成编码']):
|
||
df['组成编码'] = df['组成编码'].apply(format_compose_code)
|
||
|
||
# 转换为 JSON 列表
|
||
json_data = df.to_json(orient='records', force_ascii=False)
|
||
result_list = json.loads(json_data)
|
||
highlight_filename, most_common_page,filename = highlight_pdf_file(input_pdf_path,highlight_list)
|
||
url = f"{prefix_url}/upload/{highlight_filename}#page={most_common_page}"
|
||
if result_list:
|
||
for ins in result_list:
|
||
source_data_list = [{
|
||
"filename": filename,
|
||
"info": content,
|
||
"url": url
|
||
}]
|
||
knowledge_source_str = json.dumps(source_data_list, ensure_ascii=False)
|
||
ins["knowledge_source"] = knowledge_source_str
|
||
|
||
all_htmltojson_results.extend(result_list)
|
||
print(f" ✅ 第 {i+1} 组提取成功，包含 {len(result_list)} 条记录")
|
||
else:
|
||
print(f" ⚠️ 第 {i+1} 组提取结果为空")
|
||
|
||
except Exception as e:
|
||
print(f"❌ 处理第 {i+1} 组时发生错误: {e}")
|
||
traceback.print_exc()
|
||
continue
|
||
|
||
# 后处理步骤
|
||
if all_htmltojson_results:
|
||
all_htmltojson_results = merge_compose(all_htmltojson_results)
|
||
all_htmltojson_results = remove_fields(all_htmltojson_results, ["序号"])
|
||
|
||
# ✅ 修复：先生成实体和关系
|
||
all_entity_results, all_relation_results = generate_entities_and_relations(all_htmltojson_results)
|
||
|
||
# ✅ 修复：循环移到 generate 之后，此时 all_entity_results 才有数据
|
||
for ins in all_entity_results:
|
||
if ins["properties"].get("名称") == devicename:
|
||
ins["type"] = lower_entity_type
|
||
else:
|
||
print("⚠️ 没有提取到任何有效数据，跳过实体关系生成。")
|
||
|
||
return all_entity_results, all_relation_results
|
||
|
||
except Exception as e:
|
||
print(f"❌ get_tujiecompose_node_relation 主流程发生严重错误: {e}")
|
||
traceback.print_exc()
|
||
return [], []
|
||
|
||
|
||
if __name__ == "__main__":
|
||
file_path = '/storage01/home/hdf/project/wxkgrag/kgextract/output/output1/122-06A0014-B01001_发动机-图解目录_content_list.json'
|
||
input_pdf_path = "/storage01/home/hdf/project/wxkgrag/kgextract/know/122-06A0014-B01001_发动机-维修手册_highlighted_fdb9e9db-98d2-4ef6-bc72-52f875453f85.pdf"
|
||
output_json_path = '/storage01/home/hdf/project/wxkgrag/kgextract/output/output2/图解组成表抽取结果.json'
|
||
|
||
try:
|
||
with open(file_path, 'r', encoding='utf-8') as file:
|
||
data = json.load(file)
|
||
except Exception as e:
|
||
print(f"读取 JSON 文件时出错：{e}")
|
||
data = []
|
||
|
||
all_entity_results, all_relation_results = asyncio.run(
|
||
get_tujiecompose_node_relation(input_pdf_path, data=data)
|
||
)
|
||
|
||
length_entity = len(all_entity_results)
|
||
length_relation = len(all_relation_results)
|
||
|
||
print(f"抽取 {length_entity} 个实体，抽取 {length_relation} 个关系")
|
||
|
||
result = {
|
||
"entities": all_entity_results,
|
||
"relations": all_relation_results
|
||
}
|
||
for ins in result["entities"]:
|
||
print(ins)
|
||
print(111111111111111111111)
|
||
|
||
try:
|
||
with open(output_json_path, 'w', encoding='utf-8') as f:
|
||
json.dump(result, f, ensure_ascii=False, indent=4)
|
||
print(f"结果已保存至：{output_json_path}")
|
||
except Exception as e:
|
||
print(f"保存 JSON 文件时出错：{e}")
|
||
|
||
|