248 lines
9.0 KiB
Python
248 lines
9.0 KiB
Python
import json
|
||
import fitz # PyMuPDF
|
||
import os
|
||
import fitz
|
||
from collections import Counter
|
||
import uuid
|
||
from openai import OpenAI
|
||
from kg_build.Prompt import Prompt_compose
|
||
from extract_util import safe_json_loads,highlight_pdf_file,get_content_bbox
|
||
import re
|
||
from typing import List, Dict, Any, Tuple
|
||
from modelsAPI.model_api import OpenaiAPI
|
||
import asyncio # <-- 新增导入
|
||
|
||
import re
|
||
|
||
def is_compose_table_by_first_tr(table_body: str) -> bool:
|
||
if not isinstance(table_body, str):
|
||
return False
|
||
|
||
# 非贪婪匹配第一个 <tr> ... </tr>
|
||
match = re.search(r'<tr[^>]*>(.*?)</tr>', table_body, re.IGNORECASE | re.DOTALL)
|
||
if not match:
|
||
return False
|
||
|
||
first_tr_html = match.group(1)
|
||
|
||
# 去除所有 HTML 标签，只保留纯文本
|
||
text_in_first_tr = re.sub(r'<[^>]+>', '', first_tr_html)
|
||
|
||
# 统一转为小写便于匹配
|
||
text_lower = text_in_first_tr.lower()
|
||
|
||
# 定义多组关键词组合（任意一组满足即可）
|
||
keyword_groups = [
|
||
["序号", "组成部分", "功能", "数量"],
|
||
["序号", "组成部分", "功能"]
|
||
]
|
||
|
||
# 检查是否存在至少一组关键词全部出现在首行中
|
||
for group in keyword_groups:
|
||
if all(kw in text_lower for kw in group):
|
||
return True
|
||
|
||
return False
|
||
|
||
|
||
async def get_compose_chunk_bbox_and_removed_data(data):
|
||
"""
|
||
返回：
|
||
- 符合条件的组成表分组（每组 [前前, 前, 当前]）
|
||
- 从原始 data 中移除了所有相关记录后的新 data 列表
|
||
"""
|
||
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 is_compose_table_by_first_tr(table_body):
|
||
current_id = ins["id"]
|
||
prev_ids = [current_id - 2, current_id - 1]
|
||
|
||
# 收集要移除的 ID（前两条 + 当前）
|
||
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)
|
||
|
||
# 过滤原始 data
|
||
filtered_data = [item for item in data if item["id"] not in ids_to_remove]
|
||
|
||
return results, filtered_data
|
||
|
||
|
||
# def openai_chat(query: str, model="Qwen3-32B"):
|
||
# """调用OpenAI API"""
|
||
# try:
|
||
# client = OpenAI(
|
||
# api_key="none",
|
||
# base_url="http://192.168.0.46:59800/v1"
|
||
# )
|
||
|
||
# response = client.chat.completions.create(
|
||
# model=model,
|
||
# messages=[
|
||
# {"role": "user", "content": query},
|
||
# ],
|
||
# temperature=0.1,
|
||
# stream=False,
|
||
# response_format={"type": "json_object"},
|
||
# extra_body={"chat_template_kwargs": {"enable_thinking": False}}
|
||
# )
|
||
# return response.choices[0].message.content
|
||
# except Exception as e:
|
||
# print(f"调用OpenAI API时出错: {e}")
|
||
# return None
|
||
|
||
async def extract_compose_table_entity(chunk):
|
||
"""
|
||
提取组成表的设备、零部件实体以及关系
|
||
"""
|
||
print("提取组成表实体和关系")
|
||
|
||
# 建议：将 all_results 改为复数形式以符合习惯
|
||
|
||
|
||
final_prompt = Prompt_compose.format(text=chunk)
|
||
raw_response = await OpenaiAPI.openai_chat_aysnc(final_prompt,timeout=480)
|
||
|
||
if raw_response is None:
|
||
print(f" ⚠️ 切片 API 调用失败")
|
||
|
||
# 尝试解析 JSON
|
||
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)
|
||
|
||
# 打印美化版
|
||
# print(json.dumps(json_result, ensure_ascii=False, indent=2))
|
||
|
||
except Exception as e:
|
||
print(f" ⚠️切片 JSON 解析失败: {e}")
|
||
print(f" 原始响应: {repr(raw_response[:500])}")
|
||
json_result = []
|
||
|
||
|
||
# 返回按切片分组的结果列表
|
||
return json_result
|
||
|
||
|
||
async def get_compose_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
|
||
results, filtered_data =await get_compose_chunk_bbox_and_removed_data(data=data)
|
||
print(f"抽取的组成表总数量为：{len(results)}")
|
||
|
||
|
||
if not results:
|
||
print("未找到符合条件的组成表格，跳过处理。")
|
||
all_entity_results,all_relation_results = [],[]
|
||
else:
|
||
all_entity_results,all_relation_results = [],[] # 用于收集所有成功提取的结果
|
||
|
||
for i, compose in enumerate(results):
|
||
try:
|
||
# 获取内容和高亮信息
|
||
print(f"开始处理第{i+1}个组成表")
|
||
|
||
content, highlight_list = get_content_bbox(compose)
|
||
print("组成表切片内容为：",content)
|
||
print(len(content))
|
||
# 提取实体
|
||
result =await extract_compose_table_entity(content)
|
||
if result is None:
|
||
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 = '/storage01/home/hdf/project/wxkgrag/kgextract/output/output1/122-06A0014-B01001_发动机-维修手册_content_list.json'
|
||
# input_pdf_path = "/storage01/home/hdf/project/agent/测试.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_compose_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
|
||
# }
|
||
|
||
# # 保存为 JSON 文件
|
||
# 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}")
|
||
|