更新 chunk_text.py

修复多切片高亮问题
This commit is contained in:
Defeng 2026-07-29 18:22:54 +08:00
parent 2c5d783d2f
commit 0462e5066f

View File

@ -15,6 +15,7 @@ except ImportError:
from pathlib import PurePath, PurePosixPath, Path as PathLib
from fastapi import FastAPI, File, Path, UploadFile, HTTPException, Form, Request, Header, Body
from fileparse_util import process_document
from generate_summary import generate_summaries
from io import BytesIO
from urllib.parse import unquote, urlparse
import asyncio
@ -44,6 +45,42 @@ except ImportError:
SHIP_MODEL_NAME = ""
logger = logging.getLogger(__name__)
SUMMARY_GENERATION_TIMEOUT_SECONDS = 60
def get_text_level_summary_fallback(content_list) -> str:
summary_parts = []
for record in content_list:
if not isinstance(record, dict):
continue
text_level = record.get("text_level")
if text_level is None or text_level == "":
continue
text = record.get("text")
if text:
summary_parts.append(str(text).strip())
return "\n".join(summary_parts)
async def generate_summary_from_md(md_content: Optional[str],content_list) -> str:
if not md_content or not str(md_content).strip():
return get_text_level_summary_fallback(content_list)
try:
summaries = await asyncio.wait_for(
generate_summaries([str(md_content)]),
timeout=SUMMARY_GENERATION_TIMEOUT_SECONDS,
)
summary = summaries[0] if summaries else ""
if summary:
return summary
logger.warning("生成文档摘要为空,使用 content_list text_level 兜底摘要")
except Exception as exc:
logger.error(f"生成文档摘要失败,使用 content_list text_level 兜底摘要: {exc}", exc_info=True)
return get_text_level_summary_fallback(content_list)
ocr_engine = RapidOCR()
OCR_SEMAPHORE = asyncio.Semaphore(max(1, OCR_CONCURRENCY))
@ -263,22 +300,6 @@ def extract_image_text_sync(image_path: Union[str, PathLib]) -> dict:
}
def build_image_markdown_with_ocr(markdown_image: str, local_image_path: Optional[PathLib]) -> str:
if any(prefix in markdown_image for prefix in IMAGE_OCR_PREFIXES):
return markdown_image
if not local_image_path:
return markdown_image
try:
ocr_data = extract_image_text_sync(local_image_path)
except Exception as exc:
logger.warning(f"Failed to OCR local image {local_image_path}: {exc}")
return markdown_image
ocr_full_text = (ocr_data.get("full_text") or "").strip()
if not ocr_full_text:
return markdown_image
return f"{markdown_image}\n{ocr_full_text}"
async def get_image_ocr_text(local_image_path: Optional[PathLib]) -> str:
@ -315,6 +336,43 @@ def join_nonempty_parts(*parts: Any) -> str:
return "\n".join(part for part in (stringify_content_field(part).strip() for part in parts) if part)
def strip_image_ocr_prefix(text: Any) -> str:
text = stringify_content_field(text).strip()
for prefix in IMAGE_OCR_PREFIXES:
if text.startswith(prefix):
return text[len(prefix):].strip()
return text
def normalize_image_ocr_text(value: Any) -> str:
return "\n".join(
line
for line in (strip_image_ocr_prefix(line) for line in stringify_content_field(value).splitlines())
if line
)
def split_image_caption_and_ocr(caption: Any) -> Tuple[str, str]:
if isinstance(caption, list):
raw_items = caption
else:
raw_items = [caption]
caption_parts = []
ocr_parts = []
for item in raw_items:
for line in stringify_content_field(item).splitlines():
line = line.strip()
if not line:
continue
ocr_line = strip_image_ocr_prefix(line)
if ocr_line != line:
ocr_parts.append(ocr_line)
else:
caption_parts.append(line)
return "\n".join(caption_parts), "\n".join(ocr_parts)
def record_to_chunk_text(ins: Dict[str, Any]) -> str:
record_type = ins.get("type")
ocr_text = ins.get("ocr_text", "")
@ -328,23 +386,24 @@ def record_to_chunk_text(ins: Dict[str, Any]) -> str:
if record_type == "table":
return join_nonempty_parts(
ocr_text,
ins.get("table_caption"),
clean_table_body(ins.get("table_body", "")),
ins.get("table_footnote"),
)
if record_type == "image":
return join_nonempty_parts(
ocr_text,
ins.get("image_caption"),
ins.get("img_path"),
ins.get("image_footnote"),
)
imagecontent = ''
image_caption, caption_ocr_text = split_image_caption_and_ocr(ins.get("image_caption"))
image_caption = image_caption +"如下所示:"
img_path = stringify_content_field(ins.get("img_path")).strip()
if caption_ocr_text:
imagecontent = image_caption + "\n" + img_path + "\n" + f"备注:图片中包含的内容为{caption_ocr_text}"
else:
imagecontent = image_caption + "\n" + img_path
return imagecontent
if record_type == "chart":
return join_nonempty_parts(
ocr_text,
ins.get("chart_caption"),
ins.get("content"),
ins.get("img_path"),
@ -353,26 +412,23 @@ def record_to_chunk_text(ins: Dict[str, Any]) -> str:
if record_type == "equation":
return join_nonempty_parts(
ocr_text,
ins.get("text"),
ins.get("img_path"),
)
if record_type == "code":
return join_nonempty_parts(
ocr_text,
ins.get("code_caption"),
ins.get("code_body"),
ins.get("code_footnote"),
)
if record_type == "list":
return join_nonempty_parts(ocr_text, ins.get("list_items"))
return join_nonempty_parts(ins.get("list_items"))
if record_type in {"discarded", "header", "footer", "page_number"}:
return join_nonempty_parts(ocr_text, ins.get("text"))
return join_nonempty_parts(ins.get("text"))
return join_nonempty_parts(ocr_text, ins.get("text"))
return join_nonempty_parts(ins.get("text"))
def is_supported_chunk_record(ins: Dict[str, Any]) -> bool:
@ -393,7 +449,7 @@ def is_supported_chunk_record(ins: Dict[str, Any]) -> bool:
def append_caption_ocr(ins: Dict[str, Any], ocr_texts: List[str]) -> None:
caption_field = {
"image": "image_caption",
"image": "ocr_text",
"table": "table_caption",
"chart": "chart_caption",
"code": "code_caption",
@ -403,7 +459,9 @@ def append_caption_ocr(ins: Dict[str, Any], ocr_texts: List[str]) -> None:
caption = ins.get(caption_field, [])
existing_text = stringify_content_field(caption)
if caption_field == "ocr_text":
if ins.get("type") == "image":
existing_text = join_nonempty_parts(existing_text, ins.get("image_caption"))
elif caption_field == "ocr_text":
existing_text = join_nonempty_parts(existing_text, ins.get("text"))
additions = []
@ -772,7 +830,8 @@ def split_content_bbox(data, max_length=8000):
for ins in data:
if "type" not in ins:
continue
if not ins.get("bbox"):
is_parent_context = ins.get("_is_parent_context", False)
if not ins.get("bbox") and not is_parent_context:
continue
# 提取通用字段
bbox = ins.get("bbox")
@ -784,7 +843,8 @@ def split_content_bbox(data, max_length=8000):
new_text = record_to_chunk_text(ins)
if new_text.strip():
new_text += "\n"
should_record_highlight = True
# 父章节只作为子章节的文本上下文,不重复写入其 page_idx/bbox。
should_record_highlight = not is_parent_context
# --- 关键:即使 new_text 为空,只要 should_record_highlight 为 True就要处理 chunk 切分和 highlight 记录 ---
@ -803,7 +863,7 @@ def split_content_bbox(data, max_length=8000):
contents.append(current_content.rstrip('\n'))
highlight_lists.append(current_highlights)
current_content = new_text
current_highlights = [highlight_item]
current_highlights = [highlight_item] if should_record_highlight else []
else:
# 追加内容(如果 new_text 非空)
if new_text.strip():
@ -875,7 +935,11 @@ def group_records(records):
section_stack.pop()
parent_group = section_stack[-1][1] if section_stack else []
merged_group = parent_group + group
parent_context = [
{**record, "_is_parent_context": True}
for record in parent_group
]
merged_group = parent_context + group
result.append(merged_group)
section_stack.append((level, merged_group))
@ -1086,6 +1150,8 @@ async def data_replace(data, prefix):
使用 pathlib 安全处理路径
"""
for ins in data:
if isinstance(ins, dict) and ins.get("type") == "equation":
continue
if isinstance(ins, dict) and "img_path" in ins:
img_path = ins["img_path"]
local_image_paths = []
@ -1255,102 +1321,51 @@ 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"):
"""
合并过短的切片
- 如果某切片 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
def content_starts_with_same_lines(content: str, prefix: str) -> bool:
prefix_lines = [line.strip() for line in (prefix or "").splitlines() if line.strip()]
content_lines = [line.strip() for line in (content or "").splitlines() if line.strip()]
return bool(prefix_lines) and content_lines[:len(prefix_lines)] == prefix_lines
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
})
def content_ends_with_same_lines(content: str, suffix: str) -> bool:
suffix_lines = [line.strip() for line in (suffix or "").splitlines() if line.strip()]
content_lines = [line.strip() for line in (content or "").splitlines() if line.strip()]
return bool(suffix_lines) and content_lines[-len(suffix_lines):] == suffix_lines
# 收尾:如果末尾还有未合并的短切片(后面没有正常切片可合并)
# 则反向合并到上一个切片末尾
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("未找到匹配的舰船信息")
def is_heading_only_content(content: str) -> bool:
lines = [
line.strip()
for line in (content or "").splitlines()
if line.strip()
]
meaningful_lines = [
line
for line in lines
if not line.startswith("#文件名为:") and not re.fullmatch(r"\d+", line)
]
return bool(meaningful_lines) and all(line.startswith("#") for line in meaningful_lines)
# 安全处理 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}'
suffix = f'({info})'
# 将信息插入到第一行末尾
if newline_idx == -1:
ins['content'] = content + suffix
else:
ins['content'] = content[:newline_idx] + suffix + content[newline_idx:]
def merge_prefix_content(merged_prefix: str, content: str, pending_positions: list, positions: list):
prefix_lines = [line.strip() for line in (merged_prefix or "").splitlines() if line.strip()]
content_lines = [line.strip() for line in (content or "").splitlines() if line.strip()]
return result
duplicate_count = 0
for prefix_line, content_line in zip(prefix_lines, content_lines):
if prefix_line != content_line:
break
duplicate_count += 1
if duplicate_count == len(prefix_lines):
return content, positions
content_lines_raw = (content or "").splitlines()
content_without_duplicates = "\n".join(content_lines_raw[duplicate_count:]).lstrip("\n")
merged_content = merged_prefix + ("\n" if merged_prefix and content_without_duplicates else "") + content_without_duplicates
merged_positions = pending_positions + positions[duplicate_count:]
return merged_content, merged_positions
def merge_short_slices(slices, min_length=30, filename="122-06A0014-B01003_雷达-使用说明书.pdf"):
@ -1379,16 +1394,29 @@ def merge_short_slices(slices, min_length=30, filename="122-06A0014-B01003_雷
content = ins.get("content", "") or ""
positions = ins.get("positions", []) or []
if len(content) <= min_length:
if len(content) <= min_length or is_heading_only_content(content):
# 暂存,等到下一个正常长度的切片再合并
pending_contents.append(content)
pending_positions.extend(positions)
if pending_contents:
merged_prefix = "\n".join(pending_contents)
if content_starts_with_same_lines(content, merged_prefix):
pending_contents = [content]
pending_positions = positions
else:
pending_contents.append(content)
pending_positions.extend(positions)
else:
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
content, positions = merge_prefix_content(
merged_prefix,
content,
pending_positions,
positions,
)
pending_contents = []
pending_positions = []
@ -1403,8 +1431,9 @@ def merge_short_slices(slices, min_length=30, filename="122-06A0014-B01003_雷
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
if not content_ends_with_same_lines(last.get("content", ""), merged_suffix):
last["content"] = (last["content"] or "") + ("\n" if last["content"] else "") + merged_suffix
last["positions"] = (last["positions"] or []) + pending_positions
else:
# 极端情况:所有切片都很短,整体作为一个切片返回
result.append({
@ -1490,14 +1519,16 @@ async def process_pdf_file(pdf_path: Path, image_prefix: str,filename:str) -> Op
"""异步调用 PDF 分析服务并处理响应"""
doc_result = await process_document(filename)
if doc_result is not None:
content_list, images = doc_result
content_list, images, md_content = doc_result
if content_list and images:
replaced_content_textlevel = reset_textlevel(content_list)
replaced_content_list = await data_replace(replaced_content_textlevel, image_prefix)
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)
return {"slices": slices_check, "images": images}
summary = await generate_summary_from_md(md_content,content_list)
summary = f"文件名:{filename}\n" + summary
return {"slices": slices_check, "images": images, "summary": summary}
api_url = get_api_url()
try:
@ -1516,19 +1547,20 @@ async def process_pdf_file(pdf_path: Path, image_prefix: str,filename:str) -> Op
data = result.get("data", {})
content_list = data.get("content_list", [])
images = data.get("images", {})
md_content = data.get("full_content", "")
replaced_content_textlevel = reset_textlevel(content_list)
replaced_content_list = await data_replace(replaced_content_textlevel, image_prefix)
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}
summary = await generate_summary_from_md(md_content,content_list)
summary = f"文件名:{filename}\n" + summary
return {"slices": slices_check, "images": images, "summary": summary}
except json.JSONDecodeError:
logger.error("❌ 响应不是有效的 JSON 格式")
logger.debug(response.text)
return {"slices": [], "images": {}}
return {"slices": [], "images": {}, "summary": ""}
except Exception as e:
logger.error(f"处理 PDF 文件时出错: {e}", exc_info=True)
return None
@ -1540,13 +1572,15 @@ async def process_other_file(pdf_path: Path, image_prefix: str, filename: str) -
# 优先尝试走本地/缓存处理(保持与你原逻辑一致)
doc_result = await process_document(filename)
if doc_result is not None:
content_list, images = doc_result
content_list, images, md_content = doc_result
if content_list and images:
replaced_content_textlevel = reset_textlevel(content_list)
replaced_content_list = await data_replace(replaced_content_textlevel, image_prefix)
slices = get_chunk_bbox(replaced_content_list)
slices_check = chunk_check(slices, 8000)
return {"slices": slices_check, "images": images}
summary = await generate_summary_from_md(md_content,content_list)
summary = f"文件名:{filename}\n" + summary
return {"slices": slices_check, "images": images, "summary": summary}
api_url = get_api_url()
analyze_other_url = api_url.replace("/analyze-pdf", "/analyze-otherfile")
@ -1571,19 +1605,21 @@ async def process_other_file(pdf_path: Path, image_prefix: str, filename: str) -
data = result.get("data", {})
content_list = data.get("content_list", [])
images = data.get("images", {})
md_content = data.get("full_content", "")
# 后续的数据替换与切片处理逻辑
replaced_content_textlevel = reset_textlevel(content_list)
replaced_content_list = await data_replace(replaced_content_textlevel, image_prefix)
slices = get_chunk_bbox(replaced_content_list)
slices_check = chunk_check(slices, 8000)
return {"slices": slices_check, "images": images}
summary = await generate_summary_from_md(md_content,content_list)
summary = f"文件名:{filename}\n" + summary
return {"slices": slices_check, "images": images, "summary": summary}
except json.JSONDecodeError:
error_text = response.text if 'response' in locals() else "请求未到达服务器"
logger.error(f"❌ 响应不是有效的 JSON 格式: {error_text}")
return {"slices": [], "images": {}}
return {"slices": [], "images": {}, "summary": ""}
except Exception as e:
logger.error(f"处理 Office 文件时出错: {e}", exc_info=True)
return None
@ -1592,7 +1628,7 @@ async def process_other_file(pdf_path: Path, image_prefix: str, filename: str) -
if __name__ == "__main__":
filepath = r"E:\ZKYNLP\Hjunproject\project0506\kgrag\舰船抗沉损管训练仿真系统研究_content_list.json"
filepath = r"E:\ZKYNLP\Hjunproject\kgrag\app\信号与系统_第三版_郑君里_上_content_list.json"
try:
with open(filepath, 'r', encoding='utf-8') as file:
data = json.load(file)
@ -1609,5 +1645,10 @@ if __name__ == "__main__":
slices_2 = merge_short_slices(slices_1, min_length=30) # ← 新增这一行
for ins in slices_2[:30]:
print(ins)
for ins in slices_2:
# if any(
# position.get("page_idx") == 60
# for position in ins.get("positions", [])
# ):
print(ins)