From 0462e5066f58f6140a57a47da2f298d984585e77 Mon Sep 17 00:00:00 2001 From: Defeng Date: Wed, 29 Jul 2026 18:22:54 +0800 Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0=20chunk=5Ftext.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 修复多切片高亮问题 --- chunk_text.py | 331 ++++++++++++++++++++++++++++---------------------- 1 file changed, 186 insertions(+), 145 deletions(-) diff --git a/chunk_text.py b/chunk_text.py index 8a6239a..4e84a9f 100644 --- a/chunk_text.py +++ b/chunk_text.py @@ -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) +