339 lines
14 KiB
Python
339 lines
14 KiB
Python
import json
|
||
import fitz # PyMuPDF
|
||
import os
|
||
from collections import Counter
|
||
from kg_build.Prompt import Prompt_guzhang
|
||
from openai import OpenAI
|
||
|
||
|
||
def get_compose_chunk_bbox_and_removed_data(data):
|
||
"""
|
||
返回：
|
||
- 符合条件的故障表分组（每组 [前前, 前, 当前]）
|
||
- 从原始 data 中移除了所有相关记录后的新 data 列表
|
||
"""
|
||
title = "组成"
|
||
|
||
# 支持多个 label 模式（可扩展）
|
||
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></tr><tr><td rowspan=1 colspan=1>1</td>',
|
||
# 可在此添加更多变体，例如：
|
||
# '<table><tr><th>序号</th><th>故障现象</th><th>故障原因</th><th>维修项目</th></tr>',
|
||
# '<table><tr><td>序号</td><td>故障现象</td><td>故障原因</td><td>维修项目</td></tr>',
|
||
]
|
||
|
||
# 构建 id -> record 映射
|
||
record_map = {item["id"]: item for item in data}
|
||
|
||
results = []
|
||
ids_to_remove = set()
|
||
|
||
for ins in data:
|
||
if ins.get("type") == "table":
|
||
# 检查 caption 是否包含 "故障表"
|
||
caption_has_title = any(title in caption for caption in ins.get("table_caption", []))
|
||
|
||
# 检查 body 是否匹配任意一个 label
|
||
body = ins.get("table_body", "")
|
||
body_matches_any_label = any(label in body for label in labels)
|
||
|
||
if caption_has_title and body_matches_any_label:
|
||
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 get_tujiecompose_chunk_bbox_and_removed_data(data):
|
||
"""
|
||
返回：
|
||
- 符合条件的故障表分组（每组 [前前, 前, 当前]）
|
||
- 从原始 data 中移除了所有相关记录后的新 data 列表
|
||
"""
|
||
title = "组成"
|
||
|
||
# 支持多个 label 模式（可扩展）
|
||
labels = [
|
||
'<table><tr><td colspan=\"1\" rowspan=\"1\">组成编码</td><td colspan=\"1\" rowspan=\"1\">名称</td><td colspan=\"1\" rowspan=\"1\">是否为关重件</td><td colspan=\"1\" rowspan=\"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></tr><tr><td rowspan=1 colspan=1>1</td>'
|
||
# 可在此添加更多变体，例如：
|
||
# '<table><tr><th>序号</th><th>故障现象</th><th>故障原因</th><th>维修项目</th></tr>',
|
||
# '<table><tr><td>序号</td><td>故障现象</td><td>故障原因</td><td>维修项目</td></tr>',
|
||
]
|
||
|
||
# 构建 id -> record 映射
|
||
record_map = {item["id"]: item for item in data}
|
||
|
||
results = []
|
||
ids_to_remove = set()
|
||
|
||
for ins in data:
|
||
if ins.get("type") == "table":
|
||
# 检查 caption 是否包含 "故障表"
|
||
caption_has_title = any(title in caption for caption in ins.get("table_caption", []))
|
||
|
||
# 检查 body 是否匹配任意一个 label
|
||
body = ins.get("table_body", "")
|
||
body_matches_any_label = any(label in body for label in labels)
|
||
|
||
if caption_has_title and body_matches_any_label:
|
||
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 get_guzhang_chunk_bbox_and_removed_data(data):
|
||
title = "故障表"
|
||
label = "<table><tr><td colspan=\"1\" rowspan=\"1\">序号</td><td colspan=\"1\" rowspan=\"1\">故障现象</td><td colspan=\"1\" rowspan=\"1\">故障原因</td><td colspan=\"1\" rowspan=\"1\">维修项目</td></tr>"
|
||
|
||
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", []))
|
||
body_has_label = label in ins.get("table_body", "")
|
||
if caption_has_title and body_has_label:
|
||
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
|
||
|
||
|
||
START_MARKERS = [
|
||
# 标准格式（带 rowspan/colspan）
|
||
"<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><td rowspan=1 colspan=1>维修间隔期</td>",
|
||
|
||
# 可在此添加其他格式，例如：
|
||
# "<table><tr><td>维修项目编号</td><td>组成编码</td><td>名称</td><td>维修项目</td><td>维修间隔期</td>",
|
||
]
|
||
|
||
def extract_maintenance_groups_and_remaining(records, markers=START_MARKERS, 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:
|
||
body = record.get("table_body", "")
|
||
is_start = (record.get("type") == "table" and
|
||
any(marker in body for marker in markers))
|
||
|
||
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:
|
||
# 如果指定了 max_following_for_last，且这是最后一个组（暂时不知道是不是最后一个）
|
||
# 我们无法在遍历时知道“这是不是最后一个起始标记”
|
||
# 所以更好的办法是：先按原逻辑分组，再处理最后一个组
|
||
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]
|
||
# 更新 used_ids：移除被截掉的 ID
|
||
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
|
||
|
||
|
||
START_MARKERS_ope = [
|
||
'<table><tr><td>项目编号</td><td>操作项目</td></tr><tr>',
|
||
'<table><tr><td rowspan=1 colspan=1>项目编号</td><td rowspan=1 colspan=1>操作项目</td></tr><tr>',
|
||
'<table><tr><td colspan="1" rowspan="1">项目编号</td><td colspan="1" rowspan="1">操作项目</td></tr><tr>'
|
||
]
|
||
|
||
def extract_operation_groups_and_remaining(records, markers=START_MARKERS_ope, 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:
|
||
body = record.get("table_body", "")
|
||
is_start = (record.get("type") == "table" and
|
||
any(marker in body for marker in markers))
|
||
|
||
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:
|
||
# 如果指定了 max_following_for_last，且这是最后一个组（暂时不知道是不是最后一个）
|
||
# 我们无法在遍历时知道“这是不是最后一个起始标记”
|
||
# 所以更好的办法是：先按原逻辑分组，再处理最后一个组
|
||
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]
|
||
# 更新 used_ids：移除被截掉的 ID
|
||
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
|
||
|
||
|
||
|
||
# def get_chunk_filter(data,max_length=7000):
|
||
# for idx, item in enumerate(data, start=1):
|
||
# item["id"] = idx
|
||
# groups, remaining_records_repair = extract_maintenance_groups_and_remaining(data)
|
||
# groups, remaining_records_operation = extract_operation_groups_and_remaining(remaining_records_repair)
|
||
# groups, remaining_records_compose = get_compose_chunk_bbox_and_removed_data(remaining_records_operation)
|
||
# groups, remaining_records_composetujie =get_tujiecompose_chunk_bbox_and_removed_data(remaining_records_compose)
|
||
# results, filtered_data=get_guzhang_chunk_bbox_and_removed_data(remaining_records_composetujie)
|
||
# res = filter_content_bbox(filtered_data,max_length=max_length)
|
||
# return res
|
||
def get_chunk_filter(data,max_length=10000):
|
||
"""
|
||
将 data 中的 text/table/image 内容按字符长度分块（每块不超过 max_length），
|
||
返回 [{"content": str, "highlight_list": list}, ...]
|
||
"""
|
||
chunks = []
|
||
current_content = ""
|
||
current_highlight_list = []
|
||
|
||
def _flush_chunk():
|
||
"""内部函数：将当前内容打包进 chunks，并重置"""
|
||
if current_content or current_highlight_list:
|
||
chunks.append({
|
||
"content": current_content,
|
||
"highlight_list": current_highlight_list.copy()
|
||
})
|
||
|
||
for ins in data:
|
||
if "type" not in ins:
|
||
continue
|
||
|
||
added_content = ""
|
||
bbox = ins.get("bbox", [])
|
||
page_idx = ins.get("page_idx", -1)
|
||
|
||
if ins["type"] == "text":
|
||
text_content = ins.get("text", "")
|
||
added_content = text_content + "\n"
|
||
|
||
elif ins["type"] == "image":
|
||
image_caption = ins.get("image_caption", "")
|
||
img_path = ins.get("img_path", "")
|
||
|
||
if isinstance(image_caption, list):
|
||
image_caption = "\n".join(str(item) for item in image_caption)
|
||
elif not isinstance(image_caption, str):
|
||
image_caption = str(image_caption)
|
||
|
||
if isinstance(img_path, list):
|
||
img_path = "\n".join(str(p) for p in img_path)
|
||
elif not isinstance(img_path, str):
|
||
img_path = str(img_path)
|
||
|
||
if image_caption or img_path:
|
||
added_content = image_caption + "\n" + img_path + "\n"
|
||
else:
|
||
continue # 跳过空图像项
|
||
else:
|
||
continue
|
||
|
||
# 检查加上 added_content 后是否超限
|
||
if len(current_content) + len(added_content) > max_length:
|
||
# 如果当前块非空，先 flush
|
||
if current_content:
|
||
_flush_chunk()
|
||
current_content = ""
|
||
current_highlight_list = []
|
||
|
||
# 如果单个 added_content 超过 max_length，仍需保留（避免丢弃）
|
||
if len(added_content) > max_length:
|
||
# 强制截断或直接放入（这里选择直接放入，避免信息丢失）
|
||
current_content = added_content[:max_length]
|
||
current_highlight_list.append({"page_idx": page_idx, "bbox": bbox})
|
||
_flush_chunk()
|
||
# 剩余部分？此处简化处理：忽略超长部分（或可递归拆分，但复杂）
|
||
continue
|
||
else:
|
||
current_content = added_content
|
||
current_highlight_list.append({"page_idx": page_idx, "bbox": bbox})
|
||
else:
|
||
current_content += added_content
|
||
current_highlight_list.append({"page_idx": page_idx, "bbox": bbox})
|
||
|
||
# 处理最后一块
|
||
if current_content or current_highlight_list:
|
||
chunks.append({
|
||
"content": current_content,
|
||
"highlight_list": current_highlight_list
|
||
})
|
||
|
||
return chunks
|
||
|
||
# if __name__ == "__main__":
|
||
# file_path = '/storage01/home/hdf/project/wxkgrag/kgextract/output/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"
|
||
# try:
|
||
# with open(file_path, 'r', encoding='utf-8') as file:
|
||
# data = json.load(file)
|
||
# except Exception as e:
|
||
# print(f"读取 JSON 文件时出错：{e}")
|
||
# data = []
|
||
# res = get_chunk_filter(data=data)
|
||
# for ins in res:
|
||
# print(ins)
|
||
# print(11111111111111111)
|
||
|