517 lines
18 KiB
Python
517 lines
18 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""Handle Aliyun parsing support for parse pdf."""
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import re
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Dict, List
|
||
|
||
from alibabacloud_docmind_api20220711.client import Client as DocmindClient
|
||
from alibabacloud_tea_openapi import models as open_api_models
|
||
from alibabacloud_docmind_api20220711 import models as docmind_models
|
||
from alibabacloud_tea_util import models as util_models
|
||
|
||
# Keep parser integration steps explicit so external workflow behavior stays traceable.
|
||
ALIBABA_ACCESS_KEY_ID = os.getenv("ALIBABA_ACCESS_KEY_ID", "")
|
||
ALIBABA_ACCESS_KEY_SECRET = os.getenv("ALIBABA_ACCESS_KEY_SECRET", "")
|
||
ALIBABA_ENDPOINT = os.getenv("ALIBABA_ENDPOINT", "docmind-api.cn-hangzhou.aliyuncs.com")
|
||
|
||
# Keep parser integration steps explicit so external workflow behavior stays traceable.
|
||
MAX_CHARS = 600
|
||
OVERLAP_CHARS = 80
|
||
|
||
# Keep parser integration steps explicit so external workflow behavior stays traceable.
|
||
TOC_TITLES = {"目次", "目录"}
|
||
TITLE_SUBTYPES = {"doc_title", "para_title"}
|
||
TEXT_SUBTYPES = {"para", "none"}
|
||
FIGURE_TYPES = {"figure", "figure_name", "figure_note"}
|
||
FIGURE_SUBTYPES = {"picture", "pic_title", "pic_caption"}
|
||
|
||
|
||
# Keep parser integration steps explicit so external workflow behavior stays traceable.
|
||
def init_client() -> DocmindClient:
|
||
"""Handle init client."""
|
||
if not ALIBABA_ACCESS_KEY_ID or not ALIBABA_ACCESS_KEY_SECRET:
|
||
raise ValueError("缺少阿里云文档解析凭据，请设置 ALIBABA_ACCESS_KEY_ID 和 ALIBABA_ACCESS_KEY_SECRET")
|
||
config = open_api_models.Config(
|
||
access_key_id=ALIBABA_ACCESS_KEY_ID,
|
||
access_key_secret=ALIBABA_ACCESS_KEY_SECRET,
|
||
)
|
||
config.endpoint = ALIBABA_ENDPOINT
|
||
return DocmindClient(config)
|
||
|
||
|
||
def submit_job(client: DocmindClient, file_path: str) -> str:
|
||
"""Submit job."""
|
||
file_name = Path(file_path).name
|
||
request = docmind_models.SubmitDocParserJobAdvanceRequest(
|
||
file_url_object=open(file_path, "rb"),
|
||
file_name=file_name,
|
||
file_name_extension=Path(file_path).suffix.lstrip("."),
|
||
llm_enhancement=True,
|
||
enhancement_mode="VLM",
|
||
)
|
||
runtime = util_models.RuntimeOptions()
|
||
response = client.submit_doc_parser_job_advance(request, runtime)
|
||
return response.body.data.id
|
||
|
||
|
||
def query_status(client: DocmindClient, task_id: str) -> Dict:
|
||
"""Handle query status."""
|
||
request = docmind_models.QueryDocParserStatusRequest(id=task_id)
|
||
response = client.query_doc_parser_status(request)
|
||
return response.body.data.to_map() if response.body.data else None
|
||
|
||
|
||
def wait_for_completion(client: DocmindClient, task_id: str, poll_interval: int = 5) -> bool:
|
||
"""Wait for for completion."""
|
||
while True:
|
||
status_data = query_status(client, task_id)
|
||
if not status_data:
|
||
return False
|
||
status = status_data.get("Status", "").lower()
|
||
if status == "success":
|
||
return True
|
||
elif status == "failed":
|
||
print(f"任务失败: {status_data}")
|
||
return False
|
||
print(f"任务状态: {status}, 等待中...")
|
||
time.sleep(poll_interval)
|
||
|
||
|
||
def get_result(client: DocmindClient, task_id: str, layout_num: int = 0, layout_step_size: int = 50) -> Dict:
|
||
"""Return result."""
|
||
request = docmind_models.GetDocParserResultRequest(
|
||
id=task_id,
|
||
layout_step_size=layout_step_size,
|
||
layout_num=layout_num,
|
||
)
|
||
response = client.get_doc_parser_result(request)
|
||
return response.body.data if response.body.data else None
|
||
|
||
|
||
def collect_all_results(client: DocmindClient, task_id: str, layout_step_size: int = 50) -> List[Dict]:
|
||
"""Collect all results."""
|
||
all_layouts = []
|
||
layout_num = 0
|
||
while True:
|
||
result_data = get_result(client, task_id, layout_num, layout_step_size)
|
||
if not result_data:
|
||
break
|
||
layouts = result_data.get("layouts", [])
|
||
if not layouts:
|
||
break
|
||
all_layouts.extend(layouts)
|
||
layout_num += len(layouts)
|
||
if len(layouts) < layout_step_size:
|
||
break
|
||
return all_layouts
|
||
|
||
|
||
# Keep parser integration steps explicit so external workflow behavior stays traceable.
|
||
def normalize_text(text: str) -> str:
|
||
"""Normalize text."""
|
||
text = text.replace("\r", "\n")
|
||
text = text.replace(" ", " ")
|
||
text = re.sub(r"\n+", "\n", text)
|
||
text = re.sub(r"[ \t]+", " ", text)
|
||
return text.strip()
|
||
|
||
|
||
def get_page(layout: Dict) -> int:
|
||
"""Return page."""
|
||
return layout.get("pageNum", layout.get("pageNumber", 0))
|
||
|
||
|
||
def get_text(layout: Dict) -> str:
|
||
"""Return text."""
|
||
text = normalize_text(layout.get("text", ""))
|
||
if text:
|
||
return text
|
||
return normalize_text(layout.get("markdownContent", ""))
|
||
|
||
|
||
# Keep parser integration steps explicit so external workflow behavior stays traceable.
|
||
def is_title(layout: Dict) -> bool:
|
||
"""Return whether title."""
|
||
return layout.get("type") == "title" or layout.get("subType") in TITLE_SUBTYPES
|
||
|
||
|
||
def is_text(layout: Dict) -> bool:
|
||
"""Return whether text."""
|
||
return layout.get("type") == "text" and layout.get("subType", "none") in TEXT_SUBTYPES
|
||
|
||
|
||
def is_figure(layout: Dict) -> bool:
|
||
"""Return whether figure."""
|
||
return layout.get("type") in FIGURE_TYPES or layout.get("subType") in FIGURE_SUBTYPES
|
||
|
||
|
||
def is_table(layout: Dict) -> bool:
|
||
"""Return whether table."""
|
||
return layout.get("type") == "table"
|
||
|
||
|
||
def is_toc_layout(layout: Dict) -> bool:
|
||
"""Return whether toc layout."""
|
||
text = get_text(layout)
|
||
if text in TOC_TITLES:
|
||
return True
|
||
if get_page(layout) == 1 and re.match(r"^\d+(\.\d+)*\s+.+[.。…]{2,}\s*\d+$", text):
|
||
return True
|
||
return False
|
||
|
||
|
||
def extract_table_text(layout: Dict) -> str:
|
||
"""Extract table text."""
|
||
rows = []
|
||
for cell in layout.get("cells", []):
|
||
texts = []
|
||
for cell_layout in cell.get("layouts", []):
|
||
cell_text = normalize_text(cell_layout.get("text", ""))
|
||
if cell_text:
|
||
texts.append(cell_text)
|
||
if texts:
|
||
rows.append(" ".join(texts))
|
||
return "\n".join(rows).strip()
|
||
|
||
|
||
# Keep parser integration steps explicit so external workflow behavior stays traceable.
|
||
def build_structure_nodes(layouts: List[Dict]) -> List[Dict]:
|
||
"""Build structure nodes."""
|
||
nodes = []
|
||
for layout in layouts:
|
||
if not is_title(layout):
|
||
continue
|
||
text = get_text(layout)
|
||
if not text or text in TOC_TITLES:
|
||
continue
|
||
nodes.append(
|
||
{
|
||
"unique_id": layout.get("uniqueId"),
|
||
"page": get_page(layout),
|
||
"index": layout.get("index", 0),
|
||
"level": layout.get("level", 0),
|
||
"title": text,
|
||
"type": layout.get("type"),
|
||
"sub_type": layout.get("subType"),
|
||
}
|
||
)
|
||
return nodes
|
||
|
||
|
||
# Keep parser integration steps explicit so external workflow behavior stays traceable.
|
||
def update_section_path(section_stack: List[Dict], layout: Dict) -> List[Dict]:
|
||
"""Update section path."""
|
||
level = layout.get("level", 0)
|
||
title = get_text(layout)
|
||
while section_stack and section_stack[-1]["level"] >= level:
|
||
section_stack.pop()
|
||
section_stack.append(
|
||
{
|
||
"level": level,
|
||
"title": title,
|
||
"page": get_page(layout),
|
||
"unique_id": layout.get("uniqueId"),
|
||
}
|
||
)
|
||
return section_stack
|
||
|
||
|
||
def section_path_titles(section_stack: List[Dict]) -> List[str]:
|
||
"""Handle section path titles."""
|
||
return [item["title"] for item in section_stack]
|
||
|
||
|
||
def flush_text_block(blocks: List[Dict], semantic_blocks: List[Dict], block_id: int) -> int:
|
||
"""Handle flush text block."""
|
||
if not blocks:
|
||
return block_id
|
||
|
||
texts = [item["text"] for item in blocks if item["text"]]
|
||
merged_text = "\n".join(texts).strip()
|
||
if not merged_text:
|
||
return block_id
|
||
|
||
semantic_blocks.append(
|
||
{
|
||
"semantic_id": f"semantic-{block_id}",
|
||
"block_type": "section_text",
|
||
"page_start": min(item["page"] for item in blocks),
|
||
"page_end": max(item["page"] for item in blocks),
|
||
"section_path": blocks[0]["section_path"],
|
||
"section_level": blocks[0]["section_level"],
|
||
"section_title": blocks[0]["section_title"],
|
||
"source_ids": [item["unique_id"] for item in blocks if item.get("unique_id")],
|
||
"text": merged_text,
|
||
}
|
||
)
|
||
return block_id + 1
|
||
|
||
|
||
def build_semantic_blocks(layouts: List[Dict]) -> List[Dict]:
|
||
"""Build semantic blocks."""
|
||
semantic_blocks = []
|
||
section_stack = []
|
||
pending_text_blocks = []
|
||
block_id = 1
|
||
skip_toc_page = False
|
||
|
||
for layout in layouts:
|
||
text = get_text(layout)
|
||
page = get_page(layout)
|
||
|
||
if is_toc_layout(layout):
|
||
skip_toc_page = True
|
||
continue
|
||
if skip_toc_page and page == 1:
|
||
continue
|
||
if skip_toc_page and page != 1:
|
||
skip_toc_page = False
|
||
|
||
if is_title(layout):
|
||
block_id = flush_text_block(pending_text_blocks, semantic_blocks, block_id)
|
||
pending_text_blocks = []
|
||
section_stack = update_section_path(section_stack, layout)
|
||
continue
|
||
|
||
section_path = section_path_titles(section_stack)
|
||
section_title = section_path[-1] if section_path else "未分类"
|
||
section_level = len(section_path)
|
||
|
||
if is_table(layout):
|
||
block_id = flush_text_block(pending_text_blocks, semantic_blocks, block_id)
|
||
pending_text_blocks = []
|
||
table_text = extract_table_text(layout)
|
||
if table_text:
|
||
semantic_blocks.append(
|
||
{
|
||
"semantic_id": f"semantic-{block_id}",
|
||
"block_type": "table",
|
||
"page_start": page,
|
||
"page_end": page,
|
||
"section_path": section_path,
|
||
"section_level": section_level,
|
||
"section_title": section_title,
|
||
"source_ids": [layout.get("uniqueId")],
|
||
"text": table_text,
|
||
}
|
||
)
|
||
block_id += 1
|
||
continue
|
||
|
||
if is_figure(layout):
|
||
block_id = flush_text_block(pending_text_blocks, semantic_blocks, block_id)
|
||
pending_text_blocks = []
|
||
if text:
|
||
semantic_blocks.append(
|
||
{
|
||
"semantic_id": f"semantic-{block_id}",
|
||
"block_type": "figure",
|
||
"page_start": page,
|
||
"page_end": page,
|
||
"section_path": section_path,
|
||
"section_level": section_level,
|
||
"section_title": section_title,
|
||
"source_ids": [layout.get("uniqueId")],
|
||
"text": text,
|
||
}
|
||
)
|
||
block_id += 1
|
||
continue
|
||
|
||
if is_text(layout) and text:
|
||
pending_text_blocks.append(
|
||
{
|
||
"page": page,
|
||
"text": text,
|
||
"unique_id": layout.get("uniqueId"),
|
||
"section_path": section_path,
|
||
"section_level": section_level,
|
||
"section_title": section_title,
|
||
}
|
||
)
|
||
|
||
flush_text_block(pending_text_blocks, semantic_blocks, block_id)
|
||
return semantic_blocks
|
||
|
||
|
||
# Keep parser integration steps explicit so external workflow behavior stays traceable.
|
||
def split_text_with_overlap(text: str, max_chars: int, overlap_chars: int) -> List[str]:
|
||
"""Handle split text with overlap."""
|
||
text = text.strip()
|
||
if len(text) <= max_chars:
|
||
return [text] if text else []
|
||
|
||
parts = []
|
||
start = 0
|
||
while start < len(text):
|
||
end = min(len(text), start + max_chars)
|
||
parts.append(text[start:end].strip())
|
||
if end >= len(text):
|
||
break
|
||
start = max(0, end - overlap_chars)
|
||
return [part for part in parts if part]
|
||
|
||
|
||
def build_vector_chunks(
|
||
semantic_blocks: List[Dict],
|
||
doc_id: str,
|
||
doc_title: str,
|
||
max_chars: int,
|
||
overlap_chars: int,
|
||
) -> List[Dict]:
|
||
"""Build vector chunks."""
|
||
vector_chunks = []
|
||
chunk_index = 1
|
||
|
||
for block in semantic_blocks:
|
||
pieces = split_text_with_overlap(block["text"], max_chars, overlap_chars)
|
||
for piece_index, piece in enumerate(pieces, start=1):
|
||
if block["section_path"]:
|
||
header = f"标准：{doc_title}\n章节：{' > '.join(block['section_path'])}\n\n"
|
||
else:
|
||
header = f"标准：{doc_title}\n\n"
|
||
vector_chunks.append(
|
||
{
|
||
"doc_id": doc_id,
|
||
"doc_title": doc_title,
|
||
"chunk_id": f"chunk-{chunk_index}",
|
||
"chunk_index": chunk_index,
|
||
"semantic_id": block["semantic_id"],
|
||
"chunk_type": block["block_type"],
|
||
"piece_index": piece_index,
|
||
"page_start": block["page_start"],
|
||
"page_end": block["page_end"],
|
||
"section_path": block["section_path"],
|
||
"section_level": block["section_level"],
|
||
"section_title": block["section_title"],
|
||
"source_ids": block["source_ids"],
|
||
"text": piece,
|
||
"embedding_text": header + piece,
|
||
}
|
||
)
|
||
chunk_index += 1
|
||
|
||
return vector_chunks
|
||
|
||
|
||
def parse_pdf_to_structured_chunks(
|
||
pdf_path: str,
|
||
*,
|
||
doc_id: str,
|
||
doc_title: str,
|
||
max_chars: int = MAX_CHARS,
|
||
overlap_chars: int = OVERLAP_CHARS,
|
||
poll_interval: int = 5,
|
||
) -> Dict:
|
||
"""Parse pdf to structured chunks."""
|
||
client = init_client()
|
||
task_id = submit_job(client, pdf_path)
|
||
if not wait_for_completion(client, task_id, poll_interval):
|
||
raise RuntimeError("阿里云文档解析任务失败")
|
||
layouts = collect_all_results(client, task_id)
|
||
return convert_layouts(
|
||
layouts,
|
||
doc_id=doc_id,
|
||
doc_title=doc_title,
|
||
max_chars=max_chars,
|
||
overlap_chars=overlap_chars,
|
||
)
|
||
|
||
|
||
# Keep parser integration steps explicit so external workflow behavior stays traceable.
|
||
def convert_layouts(
|
||
layouts: List[Dict],
|
||
doc_id: str,
|
||
doc_title: str,
|
||
max_chars: int,
|
||
overlap_chars: int,
|
||
) -> Dict:
|
||
"""Handle convert layouts."""
|
||
structure_nodes = build_structure_nodes(layouts)
|
||
semantic_blocks = build_semantic_blocks(layouts)
|
||
vector_chunks = build_vector_chunks(
|
||
semantic_blocks,
|
||
doc_id=doc_id,
|
||
doc_title=doc_title,
|
||
max_chars=max_chars,
|
||
overlap_chars=overlap_chars,
|
||
)
|
||
return {
|
||
"doc_id": doc_id,
|
||
"doc_title": doc_title,
|
||
"structure_nodes": structure_nodes,
|
||
"semantic_blocks": semantic_blocks,
|
||
"vector_chunks": vector_chunks,
|
||
}
|
||
|
||
|
||
# Keep parser integration steps explicit so external workflow behavior stays traceable.
|
||
def main() -> None:
|
||
"""Run the module entrypoint."""
|
||
parser = argparse.ArgumentParser(description="阿里云文档智能解析 PDF，输出三层结构 chunks")
|
||
parser.add_argument("pdf_path", help="PDF 文件路径")
|
||
parser.add_argument("--out", default="vector_chunks.json", help="输出 JSON 文件路径")
|
||
parser.add_argument("--layouts-out", dest="layouts_output", help="输出原始 layouts JSON")
|
||
parser.add_argument("--doc-id", default="GB14747-2006", help="文档 ID")
|
||
parser.add_argument("--doc-title", default="GB 14747—2006 儿童三轮车安全要求", help="文档标题")
|
||
parser.add_argument("--max-chars", type=int, default=MAX_CHARS, help="单个检索 chunk 最大字符数")
|
||
parser.add_argument("--overlap-chars", type=int, default=OVERLAP_CHARS, help="相邻检索 chunk 重叠字符数")
|
||
parser.add_argument("--poll-interval", type=int, default=5, help="轮询间隔（秒）")
|
||
args = parser.parse_args()
|
||
|
||
pdf_path = Path(args.pdf_path).expanduser().resolve()
|
||
if not pdf_path.exists():
|
||
raise FileNotFoundError(f"PDF 文件不存在: {pdf_path}")
|
||
|
||
# Keep parser integration steps explicit so external workflow behavior stays traceable.
|
||
client = init_client()
|
||
print(f"提交任务: {pdf_path}")
|
||
task_id = submit_job(client, str(pdf_path))
|
||
print(f"任务 ID: {task_id}")
|
||
|
||
# Keep parser integration steps explicit so external workflow behavior stays traceable.
|
||
print("等待任务完成...")
|
||
if not wait_for_completion(client, task_id, args.poll_interval):
|
||
print("任务失败，退出")
|
||
return
|
||
|
||
# Keep parser integration steps explicit so external workflow behavior stays traceable.
|
||
print("获取解析结果...")
|
||
layouts = collect_all_results(client, task_id)
|
||
print(f"获取到 {len(layouts)} 个布局块")
|
||
|
||
# Keep parser integration steps explicit so external workflow behavior stays traceable.
|
||
if args.layouts_output:
|
||
layouts_path = Path(args.layouts_output).expanduser().resolve()
|
||
layouts_path.write_text(json.dumps(layouts, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
print(f"原始 layouts 已写入: {layouts_path}")
|
||
|
||
# Keep parser integration steps explicit so external workflow behavior stays traceable.
|
||
print("转换为三层结构...")
|
||
data = convert_layouts(
|
||
layouts,
|
||
doc_id=args.doc_id,
|
||
doc_title=args.doc_title,
|
||
max_chars=args.max_chars,
|
||
overlap_chars=args.overlap_chars,
|
||
)
|
||
|
||
# Keep parser integration steps explicit so external workflow behavior stays traceable.
|
||
output_path = Path(args.out).expanduser().resolve()
|
||
output_path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
|
||
print(f"结构层节点数: {len(data['structure_nodes'])}")
|
||
print(f"语义层块数: {len(data['semantic_blocks'])}")
|
||
print(f"检索层块数: {len(data['vector_chunks'])}")
|
||
print(f"输出文件: {output_path}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|