优化切片，增加章节图片信息
This commit is contained in:
parent
3ffa6704ad
commit
5a2ee98dfb
235
chunk_text.py
235
chunk_text.py
@ -342,6 +342,15 @@ def record_to_chunk_text(ins: Dict[str, Any]) -> str:
|
||||
ins.get("image_footnote"),
|
||||
)
|
||||
|
||||
if record_type == "chart":
|
||||
return join_nonempty_parts(
|
||||
ocr_text,
|
||||
ins.get("chart_caption"),
|
||||
ins.get("content"),
|
||||
ins.get("img_path"),
|
||||
ins.get("chart_footnote"),
|
||||
)
|
||||
|
||||
if record_type == "equation":
|
||||
return join_nonempty_parts(
|
||||
ocr_text,
|
||||
@ -371,6 +380,7 @@ def is_supported_chunk_record(ins: Dict[str, Any]) -> bool:
|
||||
"text",
|
||||
"table",
|
||||
"image",
|
||||
"chart",
|
||||
"equation",
|
||||
"code",
|
||||
"list",
|
||||
@ -385,6 +395,7 @@ def append_caption_ocr(ins: Dict[str, Any], ocr_texts: List[str]) -> None:
|
||||
caption_field = {
|
||||
"image": "image_caption",
|
||||
"table": "table_caption",
|
||||
"chart": "chart_caption",
|
||||
"code": "code_caption",
|
||||
}.get(ins.get("type"))
|
||||
if not caption_field:
|
||||
@ -810,85 +821,67 @@ def split_content_bbox(data, max_length=8000):
|
||||
|
||||
|
||||
|
||||
# def group_records(records):
|
||||
# result = []
|
||||
# current_group = []
|
||||
|
||||
# for i, ins in enumerate(records):
|
||||
# if 'text_level' in ins and ins['text_level']:
|
||||
# # 当遇到新的text_level时，如果current_group非空，则先将其添加到结果中
|
||||
# if current_group:
|
||||
# result.append(current_group)
|
||||
# current_group = [] # 开始新一组
|
||||
# current_group.append(ins)
|
||||
# else:
|
||||
# # 如果当前记录没有text_level或其值为空，则继续添加到当前组
|
||||
# current_group.append(ins)
|
||||
|
||||
# # 别忘了将最后一个组添加到结果中
|
||||
# if current_group:
|
||||
# result.append(current_group)
|
||||
|
||||
# return result
|
||||
|
||||
def group_records(records):
|
||||
"""
|
||||
按 text_level 分组，并为每个组合适的父级标题上下文。
|
||||
|
||||
规则：
|
||||
- text_level=N (N>1) 的组 → 在组开头注入最近的 text_level=1 到 N-1 的标题
|
||||
"""
|
||||
def _group_records_without_parent_context(records):
|
||||
result = []
|
||||
current_group = []
|
||||
|
||||
# 维护标题栈：level -> {text, page_idx, bbox}
|
||||
heading_stack = {}
|
||||
|
||||
for i, ins in enumerate(records):
|
||||
text_level = ins.get('text_level')
|
||||
|
||||
if text_level is not None and isinstance(text_level, int) and text_level > 0:
|
||||
# 更新标题栈
|
||||
heading_stack[text_level] = {
|
||||
'text': ins.get('text', ''),
|
||||
'page_idx': ins.get('page_idx', -1),
|
||||
'bbox': ins.get('bbox', []),
|
||||
'level': text_level
|
||||
}
|
||||
# 清除更低层级的标题（高层级变化后，低层级失效）
|
||||
levels_to_remove = [l for l in list(heading_stack.keys()) if l > text_level]
|
||||
for l in levels_to_remove:
|
||||
del heading_stack[l]
|
||||
|
||||
# 遇到新标题，先保存当前组
|
||||
if 'text_level' in ins and ins['text_level']:
|
||||
# 当遇到新的text_level时，如果current_group非空，则先将其添加到结果中
|
||||
if current_group:
|
||||
result.append(current_group)
|
||||
current_group = []
|
||||
|
||||
# 注入父级标题（如果当前层级 > 1）
|
||||
if text_level > 1:
|
||||
prefix_records = []
|
||||
for level in sorted(heading_stack.keys()):
|
||||
if level < text_level:
|
||||
h = heading_stack[level]
|
||||
prefix_records.append({
|
||||
'type': 'text',
|
||||
'text': h['text'],
|
||||
'text_level': level,
|
||||
'page_idx': h['page_idx'],
|
||||
'bbox': h['bbox']
|
||||
})
|
||||
current_group = prefix_records + [ins]
|
||||
current_group = [] # 开始新一组
|
||||
current_group.append(ins)
|
||||
else:
|
||||
current_group = [ins]
|
||||
else:
|
||||
# 非标题记录，继续添加到当前组
|
||||
# 如果当前记录没有text_level或其值为空，则继续添加到当前组
|
||||
current_group.append(ins)
|
||||
|
||||
# 别忘了将最后一个组添加到结果中
|
||||
if current_group:
|
||||
result.append(current_group)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _get_group_level(group):
|
||||
if not group:
|
||||
return None
|
||||
|
||||
level = group[0].get("text_level")
|
||||
if isinstance(level, int) and level > 0:
|
||||
return level
|
||||
return None
|
||||
|
||||
|
||||
def group_records(records):
|
||||
"""
|
||||
Split records by text_level, then prepend ancestor section content.
|
||||
|
||||
For example, a text_level=2 section will include the nearest text_level=1
|
||||
group content before its own content. A text_level=3 section will include
|
||||
the nearest level 1 and level 2 group content before its own content.
|
||||
"""
|
||||
base_groups = _group_records_without_parent_context(records)
|
||||
result = []
|
||||
section_stack = [] # [(level, merged_group)]
|
||||
|
||||
for group in base_groups:
|
||||
level = _get_group_level(group)
|
||||
if level is None:
|
||||
result.append(group)
|
||||
continue
|
||||
|
||||
while section_stack and section_stack[-1][0] >= level:
|
||||
section_stack.pop()
|
||||
|
||||
parent_group = section_stack[-1][1] if section_stack else []
|
||||
merged_group = parent_group + group
|
||||
result.append(merged_group)
|
||||
section_stack.append((level, merged_group))
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def get_chunk_bbox_otherfile(data, max_length=8000):
|
||||
"""
|
||||
处理没有 bbox 信息的文件数据。
|
||||
@ -1262,7 +1255,7 @@ def find_ship_info_by_hull(json_file_path, data):
|
||||
except json.JSONDecodeError:
|
||||
print("错误：JSON 文件格式不正确")
|
||||
return None
|
||||
# def merge_short_slices(slices, min_length=30,filename="122-06A0014-B01003_雷达-使用说明书.pdf"):
|
||||
def merge_short_slices(slices, min_length=30,filename="122-06A0014-B01003_雷达-使用说明书.pdf"):
|
||||
"""
|
||||
合并过短的切片：
|
||||
- 如果某切片 content 长度 <= min_length，则将其合并到下一个切片的开头
|
||||
@ -1358,110 +1351,6 @@ def find_ship_info_by_hull(json_file_path, data):
|
||||
ins['content'] = content[:newline_idx] + suffix + content[newline_idx:]
|
||||
|
||||
return result
|
||||
|
||||
def merge_short_slices(slices, min_length=30, filename="122-06A0014-B01003_雷达-使用说明书.pdf"):
|
||||
"""
|
||||
合并过短的切片：
|
||||
- 如果某切片 content 长度 <= min_length，则将其合并到下一个切片的开头
|
||||
- 若处于末尾无下一个切片，则反向合并到上一个切片末尾
|
||||
- content 用换行拼接，positions 顺序拼接
|
||||
|
||||
Args:
|
||||
slices: [{"content": str, "positions": [...]}]
|
||||
min_length: 短切片的字符长度阈值（含）
|
||||
filename: 用于实体提取的文件名
|
||||
|
||||
Returns:
|
||||
合并后的 slices 列表
|
||||
"""
|
||||
if not slices:
|
||||
return slices
|
||||
|
||||
result = []
|
||||
pending_contents = [] # 缓存等待合并到"下一个"的短切片 content
|
||||
pending_positions = [] # 缓存对应的 positions
|
||||
|
||||
for ins in slices:
|
||||
content = ins.get("content", "") or ""
|
||||
positions = ins.get("positions", []) or []
|
||||
|
||||
if len(content) <= min_length:
|
||||
# 暂存，等到下一个正常长度的切片再合并
|
||||
pending_contents.append(content)
|
||||
pending_positions.extend(positions)
|
||||
else:
|
||||
# 正常切片：把暂存的短切片合并到它的开头
|
||||
if pending_contents:
|
||||
merged_prefix = "\n".join(pending_contents)
|
||||
content = merged_prefix + ("\n" if merged_prefix else "") + content
|
||||
positions = pending_positions + positions
|
||||
pending_contents = []
|
||||
pending_positions = []
|
||||
|
||||
result.append({
|
||||
"content": content,
|
||||
"positions": positions
|
||||
})
|
||||
|
||||
# 收尾：如果末尾还有未合并的短切片（后面没有正常切片可合并）
|
||||
# 则反向合并到上一个切片末尾
|
||||
if pending_contents:
|
||||
merged_suffix = "\n".join(pending_contents)
|
||||
if result:
|
||||
last = result[-1]
|
||||
last["content"] = (last["content"] or "") + ("\n" if last["content"] else "") + merged_suffix
|
||||
last["positions"] = (last["positions"] or []) + pending_positions
|
||||
else:
|
||||
# 极端情况：所有切片都很短，整体作为一个切片返回
|
||||
result.append({
|
||||
"content": merged_suffix,
|
||||
"positions": pending_positions
|
||||
})
|
||||
|
||||
try:
|
||||
final_result = get_entity(filename)
|
||||
xinghao = find_ship_info_by_hull(SHIP_MODEL_NAME, final_result)
|
||||
if xinghao is None: # 确保 xinghao 为 None 时不会报错
|
||||
xinghao = {"model_name": "", "ship_name": ""}
|
||||
print("未找到匹配的舰船信息")
|
||||
|
||||
# 安全处理 xinghao 为 None 的情况
|
||||
model_name = ""
|
||||
if xinghao and 'model_name' in xinghao:
|
||||
model_name = xinghao['model_name']
|
||||
|
||||
other_data = format_entity_text(final_result)
|
||||
logger.info(f"文件名实体提取成功: {other_data}")
|
||||
except Exception as e:
|
||||
logger.warning(f"文件名实体提取失败: {e}")
|
||||
other_data = ""
|
||||
model_name = "" # 确保异常时 model_name 有定义
|
||||
|
||||
# 将提取的信息（如舰艇名、型号）注入到每个切片的开头
|
||||
for ins in result:
|
||||
content = ins['content']
|
||||
# 如果没有提取到有效数据，则跳过注入
|
||||
|
||||
newline_idx = content.find('\n')
|
||||
info = other_data
|
||||
if model_name:
|
||||
info += f'，型号为{model_name}'
|
||||
|
||||
# ===== 修改部分：拼接文件名 + 实体信息 =====
|
||||
# 构建前缀：#文件名为：xxx.pdf(实体信息）
|
||||
filename_prefix = f"#文件名为：{filename}"
|
||||
if info:
|
||||
filename_prefix += f"({info}）"
|
||||
# 注意：这里使用全角右括号"）"以匹配你的示例，如需半角请改为 ")"
|
||||
|
||||
# 将前缀插入到 content 的最开头
|
||||
if content:
|
||||
ins['content'] = filename_prefix + "\n" + content
|
||||
else:
|
||||
ins['content'] = filename_prefix
|
||||
# ==========================================
|
||||
|
||||
return result
|
||||
def find_and_read_content_list(directory, original_filename, encoding='utf-8'):
|
||||
"""
|
||||
根据原始文件名，在指定目录及其子目录下查找对应的 _content_list.json 文件并返回内容。
|
||||
@ -1527,6 +1416,8 @@ async def process_pdf_file(pdf_path: Path, image_prefix: str,filename:str) -> Op
|
||||
slices = get_chunk_bbox(replaced_content_list)
|
||||
slices_check = chunk_check(slices, 8000)
|
||||
slices_check = merge_short_slices(slices_check, min_length=30,filename=filename)
|
||||
print(11111111111111111111111111111111111111)
|
||||
print(len(images))
|
||||
return {"slices": slices_check, "images": images}
|
||||
|
||||
except json.JSONDecodeError:
|
||||
@ -1596,7 +1487,7 @@ async def process_other_file(pdf_path: Path, image_prefix: str, filename: str) -
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
filepath ="/app/mineru_output/舰船抗沉损管训练仿真系统研究_content_list.json"
|
||||
filepath = r"E:\ZKYNLP\Hjunproject\project0506\kgrag\舰船抗沉损管训练仿真系统研究_content_list.json"
|
||||
try:
|
||||
with open(filepath, 'r', encoding='utf-8') as file:
|
||||
data = json.load(file)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user