179 lines
8.1 KiB
Python
179 lines
8.1 KiB
Python
import json
|
||
import fitz # PyMuPDF
|
||
import os
|
||
from collections import Counter
|
||
from kg_build.Prompt import prompt_extract_beipinbeijian_node_rela
|
||
from openai import OpenAI
|
||
from extract_util import safe_json_loads, highlight_pdf_file, get_content_bbox
|
||
from modelsAPI.model_api import OpenaiAPI
|
||
|
||
# 指定 JSON 文件路径
|
||
# file_path = '/storage01/home/hdf/project/wxkgrag/kgextract/output/122-06A0014-B01001_发动机-维修手册_content_list.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 = []
|
||
|
||
# 为每条记录添加 id（从1开始）
|
||
|
||
|
||
def get_beipinbeijian_chunk_bbox_and_removed_data(data):
|
||
# 定义三种可能的表头格式
|
||
labels = [
|
||
'<table><tr><td rowspan=1 colspan=1>序</td><td rowspan=1 colspan=1>名称</td><td rowspan=1 colspan=1>型专</td><td rowspan=1 colspan=1>规格</td>',
|
||
'<table><tr><td rowspan=1 colspan=1>序号</td><td rowspan=1 colspan=1>名称</td><td rowspan=1 colspan=1>型号</td><td rowspan=1 colspan=1>规格</td>',
|
||
'<<table><tr><td rowspan=2 colspan=1>序号</td><td rowspan=2 colspan=1>名称</td><td rowspan=2 colspan=1>型号</td><td rowspan=2 colspan=1>规格</td>',
|
||
'<table><tr><td rowspan=2 colspan=1>序号</td><td rowspan=2 colspan=1>名称</td><td rowspan=2 colspan=1>型号</td><td rowspan=2 colspan=1>规格</td>',
|
||
'<table><tr><td rowspan=1 colspan=1>房</td><td rowspan=1 colspan=1>名称</td><td rowspan=1 colspan=1>型专</td><td rowspan=1 colspan=1>规格</td>',
|
||
'<table><tr><td rowspan=1 colspan=23>备品备件、附件和工具清单</td></tr><tr><td rowspan=1 colspan=9>设备备品备件清单</td><td rowspan=1 colspan=5>设备型号</td><td rowspan=1 colspan=3></td><td rowspan=1 colspan=4>设备装舰台套数</td>',
|
||
'<table><tr><td rowspan=2 colspan=1>序号</td><td rowspan=2 colspan=1>名 称</td><td rowspan=2 colspan=1>型号规格</td><td rowspan=2 colspan=1>图号或标准号</td>',
|
||
'<table><tr><td rowspan=1 colspan=1>序号</td><td rowspan=1 colspan=1>名称</td><td rowspan=1 colspan=1>型专</td><td rowspan=1 colspan=1>规格</td>'
|
||
]
|
||
|
||
record_map = {item["id"]: item for item in data}
|
||
results = []
|
||
ids_to_remove = set()
|
||
|
||
for ins in data:
|
||
if ins.get("type") == "table":
|
||
table_body = ins.get("table_body", "")
|
||
if any(label in table_body for label in labels):
|
||
current_id = ins["id"]
|
||
# 获取当前及前4个记录的 id（确保 >=1）
|
||
group_ids = [current_id - i for i in range(4, -1, -1) if current_id - i >= 1]
|
||
ids_to_remove.update(group_ids)
|
||
# 按顺序构建 group（从最早到当前）
|
||
group = [record_map[pid] for pid in group_ids if pid in record_map]
|
||
results.append(group)
|
||
|
||
filtered_data = [item for item in data if item["id"] not in ids_to_remove]
|
||
return results, filtered_data
|
||
|
||
|
||
async def extract_beipinbeijian_table_entity(chunk):
|
||
|
||
print("提取备品备件类实体和关系")
|
||
|
||
final_prompt = prompt_extract_beipinbeijian_node_rela.format(text=chunk)
|
||
raw_response = await OpenaiAPI.openai_chat_aysnc(final_prompt,timeout=600)
|
||
|
||
if raw_response is None:
|
||
print("API 调用失败，跳过当前 chunk")
|
||
return []
|
||
|
||
try:
|
||
cleaned = raw_response.strip()
|
||
if cleaned.startswith("```json"):
|
||
cleaned = cleaned[7:].lstrip()
|
||
if cleaned.endswith("```"):
|
||
cleaned = cleaned[:-3].rstrip()
|
||
|
||
json_result = safe_json_loads(cleaned)
|
||
if json_result is None:
|
||
print("JSON 解析失败: 返回 None")
|
||
return []
|
||
|
||
return json_result
|
||
|
||
except Exception as e:
|
||
print(f"⚠️ JSON 解析异常: {e}")
|
||
print(f"原始响应片段: {repr(raw_response[:500])}")
|
||
return []
|
||
|
||
|
||
|
||
|
||
async def get_beipinbeijian_node_relation(input_pdf_path,data,prefix_url = "http://192.168.0.46:59085"):
|
||
for idx, item in enumerate(data, start=1):
|
||
item["id"] = idx
|
||
# 主流程
|
||
guzhang_groups, cleaned_data = get_beipinbeijian_chunk_bbox_and_removed_data(data)
|
||
|
||
if not guzhang_groups:
|
||
print("未找到符合条件的备品备件表格，跳过处理。")
|
||
all_entity_results,all_relation_results = [],[]
|
||
else:
|
||
all_entity_results,all_relation_results = [],[] # 用于收集所有成功提取的结果
|
||
|
||
for i, compose in enumerate(guzhang_groups):
|
||
try:
|
||
print(f"开始处理第{i+1}个备品备件表")
|
||
# 获取内容和高亮信息
|
||
content, highlight_list = get_content_bbox(compose)
|
||
print("备品备件表切片内容为：",content)
|
||
|
||
# 提取实体
|
||
result =await extract_beipinbeijian_table_entity(content)
|
||
if result is []:
|
||
print(f"第 {i+1} 组：实体提取失败，跳过")
|
||
continue
|
||
|
||
# 安全访问 entities 字段
|
||
entities = result.get("entities", [])
|
||
if not isinstance(entities, list):
|
||
print(f"第 {i+1} 组：entities 字段格式异常，跳过")
|
||
continue
|
||
relationships = result.get("relationships", [])
|
||
|
||
# 为每个 entity 添加 knowledge_source
|
||
|
||
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}"
|
||
# for entity in entities:
|
||
# knowledge_source = [{
|
||
# "filename": filename,
|
||
# "info": content,
|
||
# "url": url
|
||
# }]
|
||
# # 确保 properties 存在
|
||
# if "properties" not in entity or not isinstance(entity["properties"], dict):
|
||
# entity["properties"] = {}
|
||
# entity["properties"]["knowledge_source"] = knowledge_source
|
||
for entity in entities:
|
||
# 构造原始数据对象 (List of Dict)
|
||
source_data_list = [{
|
||
"filename": filename,
|
||
"info": content,
|
||
"url": url
|
||
}]
|
||
|
||
# 2. 【关键修改】将列表序列化为 JSON 字符串
|
||
# ensure_ascii=False 确保中文内容不会被转义成 \uXXXX
|
||
knowledge_source_str = json.dumps(source_data_list, ensure_ascii=False)
|
||
|
||
# 确保 properties 存在且是字典
|
||
if "properties" not in entity or not isinstance(entity["properties"], dict):
|
||
entity["properties"] = {}
|
||
|
||
# 3. 赋值字符串而不是对象列表
|
||
entity["properties"]["knowledge_source"] = knowledge_source_str
|
||
# 更新 result 中的 entities
|
||
all_entity_results.extend(entities)
|
||
all_relation_results.extend(relationships)
|
||
|
||
except Exception as e:
|
||
print(f"处理第 {i+1} 组时发生未预期错误: {e}")
|
||
continue # 跳过当前组，继续下一轮
|
||
return all_entity_results,all_relation_results
|
||
|
||
# if __name__ == "__main__":
|
||
# file_path = '/app/files/122-06A0014-B01001_发动机-维修手册_content_list.json'
|
||
# input_pdf_path = "/app/files/122-06A0014-B01001_发动机-维修手册.pdf"
|
||
|
||
# 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 = get_guzhang_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}个关系")
|
||
|
||
|