339 lines
13 KiB
Python
339 lines
13 KiB
Python
import json
|
||
import fitz # PyMuPDF
|
||
import os
|
||
from collections import Counter
|
||
import uuid
|
||
# 假设 Prompt 和 extract_util 是你本地已有的模块
|
||
from kg_build.Prompt import prompt_extract_operation_node_rela
|
||
from kg_build.extract_util import safe_json_loads, highlight_pdf_file, get_content_bbox
|
||
import re
|
||
import asyncio
|
||
from openai import AsyncOpenAI
|
||
# from modelsAPI.model_api import OpenaiAPI
|
||
|
||
|
||
async def openai_chat_aysnc(query: str, timeout: int = 180) -> str | None:
|
||
"""异步调用 OpenAI 兼容 API（如本地部署的 Qwen、vLLM、Ollama 等）"""
|
||
client = None
|
||
try:
|
||
base_url = "http://192.168.0.46:59800/v1"
|
||
api_key = "none"
|
||
|
||
client = AsyncOpenAI(
|
||
api_key=api_key,
|
||
base_url=base_url,
|
||
timeout=timeout
|
||
)
|
||
|
||
response = await client.chat.completions.create(
|
||
model="Qwen3.5-35B-A3B",
|
||
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
|
||
finally:
|
||
if client:
|
||
await client.close()
|
||
|
||
def is_operation_table_by_first_tr(table_body: str) -> bool:
|
||
"""
|
||
判断 table_body 是否为操作项目表，依据是第一行 <tr> 中是否包含 '项目编号' 和 '操作项目'
|
||
"""
|
||
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_clean = text_in_first_tr.strip()
|
||
|
||
# 检查是否同时包含两个关键词
|
||
has_code = "项目编号" in text_clean
|
||
has_item = "操作项目" in text_clean
|
||
return has_code and has_item
|
||
|
||
|
||
def extract_operation_groups_and_remaining(records, max_following_for_last=None):
|
||
"""
|
||
通过语义判断表格是否为操作项目表（基于第一行内容），并分组。
|
||
|
||
普通起始标记：组 = [start, ..., next_start - 1]
|
||
最后一个起始标记：组 = [start, ..., start + max_following_for_last] （若指定）
|
||
"""
|
||
groups = []
|
||
current_group = None
|
||
used_ids = set()
|
||
|
||
for record in records:
|
||
is_start = (
|
||
record.get("type") == "table"
|
||
and is_operation_table_by_first_tr(record.get("table_body", ""))
|
||
)
|
||
|
||
if is_start:
|
||
if current_group is not None:
|
||
groups.append(current_group)
|
||
used_ids.update(r["id"] for r in current_group)
|
||
current_group = [record]
|
||
else:
|
||
if current_group is not None:
|
||
current_group.append(record)
|
||
|
||
# 处理最后一个组
|
||
if current_group is not None:
|
||
groups.append(current_group)
|
||
used_ids.update(r["id"] for r in current_group)
|
||
|
||
# —————— 后处理：截断最后一个组（如果需要） ——————
|
||
if groups and max_following_for_last is not None:
|
||
last_group = groups[-1]
|
||
if len(last_group) > max_following_for_last + 1: # +1 是起始记录本身
|
||
truncated = last_group[:max_following_for_last + 1]
|
||
removed_ids = {r["id"] for r in last_group[max_following_for_last + 1:]}
|
||
used_ids -= removed_ids
|
||
groups[-1] = truncated
|
||
|
||
remaining_records = [r for r in records if r["id"] not in used_ids]
|
||
return groups, remaining_records
|
||
|
||
|
||
|
||
|
||
async def extract_entity_relation(chunk):
|
||
"""
|
||
批量处理一个文档的多个切片（chunks），每个切片调用大模型结构化，并返回全部结果。
|
||
"""
|
||
final_prompt = prompt_extract_operation_node_rela.replace("text", chunk)
|
||
|
||
raw_response = await openai_chat_aysnc(final_prompt, timeout=600)
|
||
|
||
if raw_response is None:
|
||
print(f" ⚠️ 切片 API 调用失败")
|
||
return []
|
||
|
||
# 尝试解析 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)
|
||
return json_result
|
||
|
||
except Exception as e:
|
||
print(f" ⚠️ 切片 JSON 解析失败: {e}")
|
||
print(f" 原始响应: {repr(raw_response[:500])}")
|
||
return []
|
||
|
||
|
||
|
||
async def get_operation_node_relation(input_pdf_path, data, prefix_url="http://192.168.0.46:59085", device="发动机"):
|
||
# 1. 初始化数据 ID (如果输入数据还没有ID)
|
||
for idx, item in enumerate(data, start=1):
|
||
if "id" not in item:
|
||
item["id"] = idx
|
||
|
||
devicename = device
|
||
|
||
# 2. 数据预处理：获取故障组 and 清洗数据
|
||
groups1, remaining_records1 = extract_operation_groups_and_remaining(data)
|
||
groups = groups1[1:] if len(groups1) > 1 else []
|
||
|
||
all_entity_results = []
|
||
all_relation_results = []
|
||
device_operation = {
|
||
"type" : "操作程序",
|
||
"properties": {
|
||
"名称": f"{devicename}的操作程序",
|
||
}
|
||
}
|
||
|
||
if not groups:
|
||
print("未找到符合条件的操作项目组，跳过后续处理。")
|
||
else:
|
||
# 3. 遍历每个维修项目组进行抽取
|
||
for i, repair_group in enumerate(groups):
|
||
try:
|
||
print(f"开始处理第 {i+1} 个操作项目表")
|
||
|
||
# 获取内容和高亮区域
|
||
content, highlight_list = get_content_bbox(repair_group)
|
||
if not content:
|
||
print(f"第 {i+1} 组：内容为空，跳过")
|
||
continue
|
||
|
||
# ================= 重试逻辑开始 =================
|
||
result = None
|
||
max_retries = 20
|
||
retry_delay = 0.5
|
||
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 attempt in range(1, max_retries + 2):
|
||
# print(content[:500]) # 调试时可打开，避免日志过多
|
||
# 1. 调用 API
|
||
result = await extract_entity_relation(content)
|
||
|
||
# 2. 验证结果有效性
|
||
is_valid = False
|
||
if result is not None and isinstance(result, dict):
|
||
# 情况 1: 标准格式 {'entities': [...], 'relationships': [...]}
|
||
if "entities" in result and isinstance(result["entities"], list) and len(result["entities"]) > 0:
|
||
is_valid = True
|
||
# 情况 2: 兼容日志中的格式 {'result': [...]}
|
||
elif "result" in result and isinstance(result["result"], list) and len(result["result"]) > 0:
|
||
is_valid = True
|
||
|
||
# 3. 判断是否成功
|
||
if is_valid:
|
||
if attempt > 1:
|
||
print(f" ✅ 第 {attempt} 次尝试成功获取有效数据。")
|
||
break
|
||
|
||
# 4. 如果失败且还有重试机会
|
||
if attempt <= max_retries:
|
||
reason = "API 返回为空" if result is None else "返回数据格式异常或内容为空"
|
||
print(f" ⚠️ 第 {attempt} 次尝试失败：{reason}。等待 {retry_delay} 秒后重试...")
|
||
await asyncio.sleep(retry_delay)
|
||
else:
|
||
print(f" ❌ 经过 {max_retries + 1} 次尝试仍无法获取有效数据，跳过该组。")
|
||
result = None
|
||
break
|
||
# ================= 重试逻辑结束 =================
|
||
|
||
if result is None:
|
||
continue
|
||
|
||
# 构造知识来源字符串 (此处保持原有逻辑，实际项目中可能需要动态生成真实 URL)
|
||
source_data_list = [{
|
||
"filename": filename, # 建议替换为真实文件名
|
||
"info": content,
|
||
"url": url
|
||
}]
|
||
knowledge_source_str = json.dumps(source_data_list, ensure_ascii=False)
|
||
|
||
# 兼容不同返回格式
|
||
entities = result.get("entities", result.get("result", []))
|
||
relationships = result.get("relationships", [])
|
||
|
||
for entity in entities:
|
||
if "properties" not in entity or not isinstance(entity["properties"], dict):
|
||
entity["properties"] = {}
|
||
entity["properties"]["knowledge_source"] = knowledge_source_str
|
||
if entity.get("type") == "操作项目":
|
||
entity["properties"]["内容"] = content
|
||
# 构建基础关系：设备 -> 操作项目
|
||
rela = {
|
||
"type": "操作时使用",
|
||
"from_entity": f"{devicename}的操作程序",
|
||
"to_entity": entity.get("properties", {}).get("名称", "未知实体")
|
||
}
|
||
all_relation_results.append(rela)
|
||
|
||
all_entity_results.extend(entities)
|
||
all_relation_results.extend(relationships)
|
||
|
||
except Exception as e:
|
||
print(f"处理第 {i+1} 组时发生未预期错误: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
continue
|
||
|
||
|
||
# ================= 【新增】实体和关系去重逻辑 =================
|
||
print(f"去重前统计：实体数量={len(all_entity_results)}, 关系数量={len(all_relation_results)}")
|
||
|
||
# 1. 实体去重
|
||
unique_entities = []
|
||
seen_entity_keys = set()
|
||
|
||
for entity in all_entity_results:
|
||
props = entity.get("properties", {})
|
||
# 定义唯一键：优先使用 "编号"，如果没有则使用 "名称"
|
||
# 这样可以防止 OT001 出现两次
|
||
entity_id = props.get("编号") or props.get("名称")
|
||
|
||
# 如果连名称都没有，生成一个临时ID避免报错，但这种情况通常视为无效数据
|
||
if not entity_id:
|
||
entity_id = str(uuid.uuid4())
|
||
|
||
if entity_id not in seen_entity_keys:
|
||
seen_entity_keys.add(entity_id)
|
||
unique_entities.append(entity)
|
||
else:
|
||
# 可选：打印跳过的重复项
|
||
# print(f" ⚠️ 跳过重复实体: {entity_id} - {props.get('名称')}")
|
||
pass
|
||
|
||
all_entity_results = unique_entities
|
||
|
||
# 2. 关系去重
|
||
unique_relations = []
|
||
seen_relation_keys = set()
|
||
|
||
for rela in all_relation_results:
|
||
# 定义关系唯一键：类型 + 源 + 目标
|
||
r_key = (rela.get("type"), rela.get("from_entity"), rela.get("to_entity"))
|
||
|
||
if r_key not in seen_relation_keys:
|
||
seen_relation_keys.add(r_key)
|
||
unique_relations.append(rela)
|
||
|
||
all_relation_results = unique_relations
|
||
|
||
relation = {'type': '使用', 'from_entity': f'{device}的操作使用', 'to_entity': f'{device}的操作程序'}
|
||
all_relation_results.append(relation)
|
||
all_entity_results.append(device_operation)
|
||
|
||
print(f"去重后统计：实体数量={len(all_entity_results)}, 关系数量={len(all_relation_results)}")
|
||
# ========================================================
|
||
|
||
|
||
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/wxkgrag/kgextract/output/output1/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 = asyncio.run(
|
||
get_operation_node_relation(input_pdf_path, data=data, device="发动机")
|
||
)
|
||
|
||
length_entity = len(all_entity_results)
|
||
length_relation = len(all_relation_results)
|
||
for ins in all_entity_results:
|
||
print(ins)
|
||
print(11111111111111111111111)
|
||
for ins in all_relation_results:
|
||
print(ins)
|
||
print(22222222222222222222)
|
||
print(f"最终抽取结果：{length_entity}个实体，{length_relation}个关系")
|
||
|
||
|