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 = [ '', # 可在此添加更多变体，例如： # '
序号组成部分功能数量
1
', # '
序号故障现象故障原因维修项目
', ] # 构建 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 = [ '
序号故障现象故障原因维修项目
', '
组成编码名称是否为关重件是否为寿命件
' # 可在此添加更多变体，例如： # '
序号组成部分功能数量
1
', # '
序号故障现象故障原因维修项目
', ] # 构建 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 = "
序号故障现象故障原因维修项目
" 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） "
序号故障现象故障原因维修项目
", # 可在此添加其他格式，例如： # "
维修项目编号组成编码名称维修项目维修间隔期
", ] 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 = [ '
维修项目编号组成编码名称维修项目维修间隔期
', '
项目编号操作项目
', '
项目编号操作项目
' ] 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)
项目编号操作项目