更新 chunk_text.py
切片优化
This commit is contained in:
parent
f8c99dcfdc
commit
715829d763
189
chunk_text.py
189
chunk_text.py
@ -65,21 +65,36 @@ def get_text_level_summary_fallback(content_list) -> str:
|
||||
return "\n".join(summary_parts)
|
||||
|
||||
|
||||
async def generate_summary_from_md(md_content: Optional[str],content_list) -> str:
|
||||
async def generate_summary_from_md(
|
||||
md_content: Optional[str],
|
||||
content_list,
|
||||
) -> str:
|
||||
fallback_summary = get_text_level_summary_fallback(content_list)
|
||||
|
||||
if not md_content or not str(md_content).strip():
|
||||
return get_text_level_summary_fallback(content_list)
|
||||
return fallback_summary
|
||||
|
||||
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 兜底摘要")
|
||||
if summary and str(summary).strip():
|
||||
return str(summary).strip()
|
||||
|
||||
logger.warning(
|
||||
"Generated summary is empty; using text_level fallback"
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error(f"生成文档摘要失败，使用 content_list text_level 兜底摘要: {exc}", exc_info=True)
|
||||
return get_text_level_summary_fallback(content_list)
|
||||
logger.error(
|
||||
"Summary generation failed; using text_level fallback: %s",
|
||||
exc,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
logger.info("Fallback summary: %s", fallback_summary)
|
||||
return fallback_summary
|
||||
|
||||
|
||||
ocr_engine = RapidOCR()
|
||||
@ -300,22 +315,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:
|
||||
@ -352,6 +351,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", "")
|
||||
@ -365,23 +401,25 @@ 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"),
|
||||
@ -390,25 +428,23 @@ def record_to_chunk_text(ins: Dict[str, Any]) -> str:
|
||||
|
||||
if record_type == "equation":
|
||||
return join_nonempty_parts(
|
||||
ocr_text,
|
||||
ins.get("text"),
|
||||
)
|
||||
|
||||
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:
|
||||
@ -429,7 +465,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",
|
||||
@ -439,7 +475,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 = []
|
||||
@ -1294,6 +1332,52 @@ def find_ship_info_by_hull(json_file_path, data):
|
||||
print("错误：JSON 文件格式不正确")
|
||||
return None
|
||||
|
||||
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
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
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()]
|
||||
|
||||
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"):
|
||||
"""
|
||||
合并过短的切片：
|
||||
@ -1320,16 +1404,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):
|
||||
# 暂存，等到下一个正常长度的切片再合并
|
||||
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 = []
|
||||
|
||||
@ -1344,6 +1441,7 @@ def merge_short_slices(slices, min_length=30, filename="122-06A0014-B01003_雷
|
||||
merged_suffix = "\n".join(pending_contents)
|
||||
if result:
|
||||
last = result[-1]
|
||||
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:
|
||||
@ -1432,7 +1530,7 @@ async def process_pdf_file(pdf_path: Path, image_prefix: str,filename:str) -> Op
|
||||
doc_result = await process_document(filename)
|
||||
if doc_result is not None:
|
||||
content_list, images, md_content = doc_result
|
||||
if content_list and images:
|
||||
if content_list:
|
||||
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)
|
||||
@ -1440,7 +1538,8 @@ async def process_pdf_file(pdf_path: Path, image_prefix: str,filename:str) -> Op
|
||||
slices_check = merge_short_slices(slices_check, min_length=30,filename=filename)
|
||||
summary = await generate_summary_from_md(md_content,content_list)
|
||||
summary = f"文件名:{filename}\n" + summary
|
||||
return {"slices": slices_check, "images": images, "summary": summary}
|
||||
logger.info(f"摘要生成:" + summary[:30])
|
||||
return {"slices": slices_check, "images": images or {}, "summary": summary}
|
||||
|
||||
api_url = get_api_url()
|
||||
try:
|
||||
@ -1540,7 +1639,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\project0506\kgrag\163-06A0014-B01001_发动机-维修手册_content_list.json"
|
||||
try:
|
||||
with open(filepath, 'r', encoding='utf-8') as file:
|
||||
data = json.load(file)
|
||||
@ -1557,5 +1656,5 @@ if __name__ == "__main__":
|
||||
slices_2 = merge_short_slices(slices_1, min_length=30) # ← 新增这一行
|
||||
|
||||
|
||||
for ins in slices_2[:30]:
|
||||
for ins in slices_2[:15]:
|
||||
print(ins)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user