Refactor document handling and update Milvus collection settings
- Removed multiple failed document entries from `documents.json`. - Added a new document entry with updated metadata and changed the index name to `regulations_dense_1024_v2`. - Updated architecture documentation to reflect changes in the Milvus collection name. - Adjusted requirements by removing the sqlalchemy dependency. - Modified test cases to align with new document structure and naming conventions. - Introduced a new test file for Milvus vector index runtime recovery and error handling. - Updated assertions in various test files to ensure compatibility with the new schema.
This commit is contained in:
ash66
8
backend/aliyun_parser/.claude/settings.local.json
Normal file
8
backend/aliyun_parser/.claude/settings.local.json
Normal file
@@ -0,0 +1,8 @@
|
||||
{
|
||||
"permissions": {
|
||||
"allow": [
|
||||
"Bash(python3 *)",
|
||||
"Bash(PGPASSWORD=postgresql123456 psql *)"
|
||||
]
|
||||
}
|
||||
}
|
||||
475
backend/aliyun_parser/parse_pdf.py
Normal file
475
backend/aliyun_parser/parse_pdf.py
Normal file
@@ -0,0 +1,475 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
阿里云文档智能 API 解析 PDF，输出三层结构 chunks
|
||||
- structure_nodes: 目录树结构
|
||||
- semantic_blocks: 语义块（章节文本、表格、图片）
|
||||
- vector_chunks: 检索块（带 overlap 切分）
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
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
|
||||
|
||||
# ===================== 阿里云配置 =====================
|
||||
ALIBABA_ACCESS_KEY_ID = "LTAI5t6fWvAsvZkoF9WTbtys"
|
||||
ALIBABA_ACCESS_KEY_SECRET = "WX4oaE4FLYRa5L85TMQkqRPHeTJAF0"
|
||||
ALIBABA_ENDPOINT = "docmind-api.cn-hangzhou.aliyuncs.com"
|
||||
|
||||
# ===================== 切分参数 =====================
|
||||
MAX_CHARS = 600
|
||||
OVERLAP_CHARS = 80
|
||||
|
||||
# ===================== 布局类型常量 =====================
|
||||
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"}
|
||||
|
||||
|
||||
# ===================== 阿里云 API 客户端 =====================
|
||||
def init_client() -> DocmindClient:
|
||||
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:
|
||||
"""提交文档解析任务"""
|
||||
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:
|
||||
"""查询任务状态"""
|
||||
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:
|
||||
"""等待任务完成"""
|
||||
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:
|
||||
"""获取解析结果"""
|
||||
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]:
|
||||
"""收集所有解析结果"""
|
||||
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
|
||||
|
||||
|
||||
# ===================== 文本处理 =====================
|
||||
def normalize_text(text: str) -> str:
|
||||
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 layout.get("pageNum", layout.get("pageNumber", 0))
|
||||
|
||||
|
||||
def get_text(layout: Dict) -> str:
|
||||
text = normalize_text(layout.get("text", ""))
|
||||
if text:
|
||||
return text
|
||||
return normalize_text(layout.get("markdownContent", ""))
|
||||
|
||||
|
||||
# ===================== 布局类型判断 =====================
|
||||
def is_title(layout: Dict) -> bool:
|
||||
return layout.get("type") == "title" or layout.get("subType") in TITLE_SUBTYPES
|
||||
|
||||
|
||||
def is_text(layout: Dict) -> bool:
|
||||
return layout.get("type") == "text" and layout.get("subType", "none") in TEXT_SUBTYPES
|
||||
|
||||
|
||||
def is_figure(layout: Dict) -> bool:
|
||||
return layout.get("type") in FIGURE_TYPES or layout.get("subType") in FIGURE_SUBTYPES
|
||||
|
||||
|
||||
def is_table(layout: Dict) -> bool:
|
||||
return layout.get("type") == "table"
|
||||
|
||||
|
||||
def is_toc_layout(layout: Dict) -> bool:
|
||||
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:
|
||||
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()
|
||||
|
||||
|
||||
# ===================== 结构层：目录树 =====================
|
||||
def build_structure_nodes(layouts: List[Dict]) -> List[Dict]:
|
||||
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
|
||||
|
||||
|
||||
# ===================== 语义层：章节内容 =====================
|
||||
def update_section_path(section_stack: List[Dict], layout: Dict) -> List[Dict]:
|
||||
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]:
|
||||
return [item["title"] for item in section_stack]
|
||||
|
||||
|
||||
def flush_text_block(blocks: List[Dict], semantic_blocks: List[Dict], block_id: int) -> int:
|
||||
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]:
|
||||
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
|
||||
|
||||
|
||||
# ===================== 检索层：向量 chunks =====================
|
||||
def split_text_with_overlap(text: str, max_chars: int, overlap_chars: int) -> List[str]:
|
||||
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]:
|
||||
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 convert_layouts(
|
||||
layouts: List[Dict],
|
||||
doc_id: str,
|
||||
doc_title: str,
|
||||
max_chars: int,
|
||||
overlap_chars: int,
|
||||
) -> Dict:
|
||||
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,
|
||||
}
|
||||
|
||||
|
||||
# ===================== CLI 入口 =====================
|
||||
def main() -> None:
|
||||
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}")
|
||||
|
||||
# 1. 提交阿里云任务
|
||||
client = init_client()
|
||||
print(f"提交任务: {pdf_path}")
|
||||
task_id = submit_job(client, str(pdf_path))
|
||||
print(f"任务 ID: {task_id}")
|
||||
|
||||
# 2. 等待完成
|
||||
print("等待任务完成...")
|
||||
if not wait_for_completion(client, task_id, args.poll_interval):
|
||||
print("任务失败，退出")
|
||||
return
|
||||
|
||||
# 3. 获取 layouts
|
||||
print("获取解析结果...")
|
||||
layouts = collect_all_results(client, task_id)
|
||||
print(f"获取到 {len(layouts)} 个布局块")
|
||||
|
||||
# 4. 输出原始 layouts（可选）
|
||||
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}")
|
||||
|
||||
# 5. 转换为三层结构
|
||||
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,
|
||||
)
|
||||
|
||||
# 6. 输出结果
|
||||
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()
|
||||
115
backend/aliyun_parser/rebuild_milvus_collection.py
Normal file
115
backend/aliyun_parser/rebuild_milvus_collection.py
Normal file
@@ -0,0 +1,115 @@
|
||||
"""Rebuild the migrated Milvus collection from saved vector chunks."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from pymilvus import Collection, CollectionSchema, DataType, FieldSchema, connections, utility
|
||||
|
||||
|
||||
DEFAULT_COLLECTION = "regulations_dense_1024_v2"
|
||||
DEFAULT_DIM = 1024
|
||||
|
||||
|
||||
def build_collection(name: str, dim: int) -> Collection:
|
||||
"""Create the migrated Milvus collection from scratch."""
|
||||
if utility.has_collection(name):
|
||||
utility.drop_collection(name)
|
||||
|
||||
schema = CollectionSchema(
|
||||
fields=[
|
||||
FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=128, is_primary=True, auto_id=False),
|
||||
FieldSchema(name="doc_id", dtype=DataType.VARCHAR, max_length=64),
|
||||
FieldSchema(name="doc_title", dtype=DataType.VARCHAR, max_length=256),
|
||||
FieldSchema(name="chunk_id", dtype=DataType.VARCHAR, max_length=128),
|
||||
FieldSchema(name="chunk_index", dtype=DataType.INT64),
|
||||
FieldSchema(name="piece_index", dtype=DataType.INT64),
|
||||
FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535),
|
||||
FieldSchema(name="embedding_text", dtype=DataType.VARCHAR, max_length=65535),
|
||||
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
||||
FieldSchema(name="semantic_id", dtype=DataType.VARCHAR, max_length=128),
|
||||
FieldSchema(name="chunk_type", dtype=DataType.VARCHAR, max_length=64),
|
||||
FieldSchema(name="page_start", dtype=DataType.INT64),
|
||||
FieldSchema(name="page_end", dtype=DataType.INT64),
|
||||
FieldSchema(name="section_level", dtype=DataType.INT64),
|
||||
FieldSchema(name="source_ids", dtype=DataType.VARCHAR, max_length=4096),
|
||||
FieldSchema(name="section_path", dtype=DataType.VARCHAR, max_length=4096),
|
||||
FieldSchema(name="section_title", dtype=DataType.VARCHAR, max_length=512),
|
||||
FieldSchema(name="metadata_json", dtype=DataType.VARCHAR, max_length=65535),
|
||||
FieldSchema(name="created_at", dtype=DataType.INT64),
|
||||
],
|
||||
description="Dense-only regulations index",
|
||||
enable_dynamic_field=False,
|
||||
)
|
||||
collection = Collection(name=name, schema=schema)
|
||||
collection.create_index(
|
||||
field_name="embedding",
|
||||
index_params={
|
||||
"metric_type": "COSINE",
|
||||
"index_type": "IVF_FLAT",
|
||||
"params": {"nlist": 128},
|
||||
},
|
||||
)
|
||||
return collection
|
||||
|
||||
|
||||
def load_chunks(payload_path: Path) -> list[dict]:
|
||||
"""Load vector chunks emitted by the Aliyun parser pipeline."""
|
||||
payload = json.loads(payload_path.read_text(encoding="utf-8"))
|
||||
if isinstance(payload, dict):
|
||||
chunks = payload.get("vector_chunks", [])
|
||||
else:
|
||||
chunks = payload
|
||||
if not isinstance(chunks, list):
|
||||
raise ValueError("vector chunk payload must be a list or a dict containing vector_chunks")
|
||||
return chunks
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Rebuild the target collection from a vector chunk payload."""
|
||||
parser = argparse.ArgumentParser(description="Rebuild the migrated Milvus collection.")
|
||||
parser.add_argument("--host", default="127.0.0.1", help="Milvus host")
|
||||
parser.add_argument("--port", default="19530", help="Milvus port")
|
||||
parser.add_argument("--collection", default=DEFAULT_COLLECTION, help="Milvus collection name")
|
||||
parser.add_argument("--dim", type=int, default=DEFAULT_DIM, help="Embedding dimension")
|
||||
parser.add_argument("--payload", required=True, help="Path to vector_chunks.json or a compatible JSON file")
|
||||
args = parser.parse_args()
|
||||
|
||||
connections.connect("default", host=args.host, port=args.port)
|
||||
collection = build_collection(args.collection, args.dim)
|
||||
chunks = load_chunks(Path(args.payload))
|
||||
if not chunks:
|
||||
print("No vector chunks found; collection was created but remains empty.")
|
||||
return
|
||||
|
||||
data = [
|
||||
[chunk["chunk_id"] for chunk in chunks],
|
||||
[chunk["doc_id"] for chunk in chunks],
|
||||
[chunk["doc_title"] for chunk in chunks],
|
||||
[chunk["chunk_id"] for chunk in chunks],
|
||||
[int(chunk.get("chunk_index", 0) or 0) for chunk in chunks],
|
||||
[int(chunk.get("piece_index", 0) or 0) for chunk in chunks],
|
||||
[str(chunk.get("text", ""))[:65535] for chunk in chunks],
|
||||
[str(chunk.get("embedding_text", chunk.get("text", "")))[:65535] for chunk in chunks],
|
||||
[chunk["embedding"] for chunk in chunks],
|
||||
[str(chunk.get("semantic_id", "")) for chunk in chunks],
|
||||
[str(chunk.get("chunk_type", "")) for chunk in chunks],
|
||||
[int(chunk.get("page_start", 0) or 0) for chunk in chunks],
|
||||
[int(chunk.get("page_end", 0) or 0) for chunk in chunks],
|
||||
[int(chunk.get("section_level", 0) or 0) for chunk in chunks],
|
||||
[json.dumps(chunk.get("source_ids", []), ensure_ascii=False) for chunk in chunks],
|
||||
[json.dumps(chunk.get("section_path", []), ensure_ascii=False) for chunk in chunks],
|
||||
[str(chunk.get("section_title", "")) for chunk in chunks],
|
||||
[json.dumps(chunk, ensure_ascii=False) for chunk in chunks],
|
||||
[int(chunk.get("created_at", 0) or 0) for chunk in chunks],
|
||||
]
|
||||
collection.insert(data)
|
||||
collection.flush()
|
||||
collection.load()
|
||||
print(f"Rebuilt collection {args.collection} with {len(chunks)} chunks.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
122
backend/aliyun_parser/schema.sql
Normal file
122
backend/aliyun_parser/schema.sql
Normal file
@@ -0,0 +1,122 @@
|
||||
-- 法规文档向量检索系统数据库表结构
|
||||
-- PostgreSQL
|
||||
|
||||
-- ==================== 文档表 ====================
|
||||
CREATE TABLE documents (
|
||||
id SERIAL PRIMARY KEY,
|
||||
doc_id VARCHAR(128) UNIQUE NOT NULL, -- 文档唯一标识，如 "GB14747-2006"
|
||||
title VARCHAR(512) NOT NULL, -- 文档标题
|
||||
doc_type VARCHAR(32), -- 文档类型：标准/法规/规范
|
||||
standard_number VARCHAR(64), -- 标准编号：如 "GB 14747-2006"
|
||||
publish_date DATE, -- 发布日期
|
||||
implement_date DATE, -- 实施日期
|
||||
status VARCHAR(32), -- 状态：现行/废止/修订
|
||||
source_url VARCHAR(512), -- 来源 URL
|
||||
file_path VARCHAR(512), -- 本地 PDF 文件路径
|
||||
file_size INT, -- 文件大小（字节）
|
||||
upload_time TIMESTAMP DEFAULT NOW(), -- 上传时间
|
||||
created_at TIMESTAMP DEFAULT NOW(),
|
||||
updated_at TIMESTAMP DEFAULT NOW()
|
||||
);
|
||||
|
||||
COMMENT ON TABLE documents IS '文档元数据表';
|
||||
COMMENT ON COLUMN documents.doc_id IS '文档唯一标识，用于关联 Milvus 和其他表';
|
||||
COMMENT ON COLUMN documents.standard_number IS '标准编号，如 GB 14747-2006';
|
||||
|
||||
-- ==================== 章节结构表 ====================
|
||||
CREATE TABLE sections (
|
||||
id SERIAL PRIMARY KEY,
|
||||
doc_id VARCHAR(128) NOT NULL,
|
||||
unique_id VARCHAR(64) NOT NULL, -- 阿里云返回的唯一标识
|
||||
level INT NOT NULL, -- 层级：1, 2, 3...
|
||||
title VARCHAR(512) NOT NULL, -- 章节标题
|
||||
page INT, -- 所在页码
|
||||
index INT, -- 页内顺序
|
||||
parent_id INT, -- 父章节 ID（树形结构）
|
||||
created_at TIMESTAMP DEFAULT NOW(),
|
||||
|
||||
CONSTRAINT fk_sections_doc_id FOREIGN KEY (doc_id) REFERENCES documents(doc_id),
|
||||
CONSTRAINT fk_sections_parent_id FOREIGN KEY (parent_id) REFERENCES sections(id),
|
||||
CONSTRAINT uq_sections_doc_unique UNIQUE (doc_id, unique_id)
|
||||
);
|
||||
|
||||
COMMENT ON TABLE sections IS '章节结构表，用于目录导航';
|
||||
COMMENT ON COLUMN sections.parent_id IS '父章节 ID，构建树形结构';
|
||||
COMMENT ON COLUMN sections.level IS '层级深度，1 为最顶层';
|
||||
|
||||
-- ==================== 语义块表 ====================
|
||||
CREATE TABLE semantic_blocks (
|
||||
id SERIAL PRIMARY KEY,
|
||||
doc_id VARCHAR(128) NOT NULL,
|
||||
semantic_id VARCHAR(64) NOT NULL, -- 语义块唯一标识
|
||||
block_type VARCHAR(32) NOT NULL, -- 类型：section_text/table/figure
|
||||
page_start INT NOT NULL, -- 起始页码
|
||||
page_end INT NOT NULL, -- 结束页码
|
||||
section_id INT, -- 所属章节
|
||||
section_title VARCHAR(512), -- 章节标题（冗余，方便查询）
|
||||
section_level INT, -- 章节层级
|
||||
source_ids JSONB, -- 原始 layout IDs（JSON 数组）
|
||||
text TEXT NOT NULL, -- 完整内容（未被切分）
|
||||
created_at TIMESTAMP DEFAULT NOW(),
|
||||
|
||||
CONSTRAINT fk_semantic_blocks_doc_id FOREIGN KEY (doc_id) REFERENCES documents(doc_id),
|
||||
CONSTRAINT fk_semantic_blocks_section_id FOREIGN KEY (section_id) REFERENCES sections(id),
|
||||
CONSTRAINT uq_semantic_blocks_doc_semantic UNIQUE (doc_id, semantic_id)
|
||||
);
|
||||
|
||||
COMMENT ON TABLE semantic_blocks IS '语义块表，用于邻域扩展，恢复完整内容';
|
||||
COMMENT ON COLUMN semantic_blocks.block_type IS '类型：section_text（正文）、table（表格）、figure（图示）';
|
||||
COMMENT ON COLUMN semantic_blocks.source_ids IS '原始阿里云 layout 的 uniqueId 数组';
|
||||
COMMENT ON COLUMN semantic_blocks.text IS '完整语义内容，未被切分';
|
||||
|
||||
-- ==================== 向量块元数据表 ====================
|
||||
CREATE TABLE vector_chunks (
|
||||
id SERIAL PRIMARY KEY,
|
||||
doc_id VARCHAR(128) NOT NULL,
|
||||
chunk_id VARCHAR(64) NOT NULL, -- Milvus 主键
|
||||
semantic_id VARCHAR(64) NOT NULL, -- 关联语义块
|
||||
chunk_index INT NOT NULL, -- 切片序号（全局）
|
||||
piece_index INT, -- 同语义块内的切片序号
|
||||
page_start INT,
|
||||
page_end INT,
|
||||
section_title VARCHAR(512),
|
||||
text VARCHAR(2048), -- 切片文本（可选，缩短版用于展示）
|
||||
source_ids JSONB, -- 原始 layout IDs（JSON 数组）
|
||||
created_at TIMESTAMP DEFAULT NOW(),
|
||||
|
||||
CONSTRAINT fk_vector_chunks_doc_id FOREIGN KEY (doc_id) REFERENCES documents(doc_id),
|
||||
CONSTRAINT fk_vector_chunks_semantic_id FOREIGN KEY (doc_id, semantic_id)
|
||||
REFERENCES semantic_blocks(doc_id, semantic_id),
|
||||
CONSTRAINT uq_vector_chunks_doc_chunk UNIQUE (doc_id, chunk_id)
|
||||
);
|
||||
|
||||
COMMENT ON TABLE vector_chunks IS '向量块元数据表，用于快速关联查询';
|
||||
COMMENT ON COLUMN vector_chunks.chunk_id IS 'Milvus 向量库主键';
|
||||
COMMENT ON COLUMN vector_chunks.piece_index IS '同语义块内的切片序号，用于按序拼接';
|
||||
|
||||
-- ==================== 索引 ====================
|
||||
CREATE INDEX idx_sections_doc_id ON sections(doc_id);
|
||||
CREATE INDEX idx_sections_parent_id ON sections(parent_id);
|
||||
CREATE INDEX idx_sections_level ON sections(level);
|
||||
|
||||
CREATE INDEX idx_semantic_blocks_doc_id ON semantic_blocks(doc_id);
|
||||
CREATE INDEX idx_semantic_blocks_section_id ON semantic_blocks(section_id);
|
||||
CREATE INDEX idx_semantic_blocks_block_type ON semantic_blocks(block_type);
|
||||
CREATE INDEX idx_semantic_blocks_semantic_id ON semantic_blocks(semantic_id);
|
||||
|
||||
CREATE INDEX idx_vector_chunks_doc_id ON vector_chunks(doc_id);
|
||||
CREATE INDEX idx_vector_chunks_semantic_id ON vector_chunks(semantic_id);
|
||||
CREATE INDEX idx_vector_chunks_chunk_id ON vector_chunks(chunk_id);
|
||||
|
||||
-- ==================== 触发器：自动更新 updated_at ====================
|
||||
CREATE OR REPLACE FUNCTION update_updated_at()
|
||||
RETURNS TRIGGER AS $$
|
||||
BEGIN
|
||||
NEW.updated_at = NOW();
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$ LANGUAGE plpgsql;
|
||||
|
||||
CREATE TRIGGER tr_documents_updated_at
|
||||
BEFORE UPDATE ON documents
|
||||
FOR EACH ROW EXECUTE FUNCTION update_updated_at();
|
||||
327
backend/aliyun_parser/upload_to_milvus.py
Normal file
327
backend/aliyun_parser/upload_to_milvus.py
Normal file
@@ -0,0 +1,327 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
将 vector_chunks.json 向量化并上传到 Milvus 和 PostgreSQL
|
||||
使用中转站的 OpenAI 兼容 API
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import List, Dict
|
||||
|
||||
import psycopg2
|
||||
from psycopg2.extras import execute_values
|
||||
from pymilvus import (
|
||||
connections,
|
||||
Collection,
|
||||
FieldSchema,
|
||||
CollectionSchema,
|
||||
DataType,
|
||||
utility,
|
||||
)
|
||||
from openai import OpenAI
|
||||
|
||||
# ===================== 配置 =====================
|
||||
# 中转站配置
|
||||
RELAY_BASE_URL = "http://6.86.80.4:30080/v1"
|
||||
RELAY_API_KEY = "sk-5HeY7gfSIlyZMacfuXOf5cphpymsNqufEu1ou4U3avbULcyY"
|
||||
EMBEDDING_MODEL = "text-embedding-v3" # 中转站支持的 embedding 模型
|
||||
|
||||
# Milvus 配置
|
||||
MILVUS_HOST = "localhost"
|
||||
MILVUS_PORT = "19530"
|
||||
COLLECTION_NAME = "regulation_chunks"
|
||||
|
||||
# PostgreSQL 配置
|
||||
PG_HOST = "6.86.80.10"
|
||||
PG_PORT = 5432
|
||||
PG_USER = "postgresql"
|
||||
PG_PASSWORD = "postgresql123456"
|
||||
PG_DATABASE = "postgres"
|
||||
|
||||
|
||||
# ===================== Embedding =====================
|
||||
def get_openai_client(api_key: str, base_url: str) -> OpenAI:
|
||||
"""创建 OpenAI 客户端连接到中转站"""
|
||||
return OpenAI(api_key=api_key, base_url=base_url)
|
||||
|
||||
|
||||
def get_embeddings_batch(client: OpenAI, texts: List[str], batch_size: int = 10) -> List[List[float]]:
|
||||
"""批量获取文本向量"""
|
||||
all_embeddings = []
|
||||
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i:i + batch_size]
|
||||
print(f"Embedding batch {i // batch_size + 1}/{(len(texts) - 1) // batch_size + 1}...")
|
||||
|
||||
response = client.embeddings.create(
|
||||
model=EMBEDDING_MODEL,
|
||||
input=batch,
|
||||
)
|
||||
|
||||
embeddings = [item.embedding for item in response.data]
|
||||
all_embeddings.extend(embeddings)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
|
||||
# ===================== Milvus =====================
|
||||
def init_milvus(host: str, port: str):
|
||||
connections.connect("default", host=host, port=port)
|
||||
print(f"已连接 Milvus: {host}:{port}")
|
||||
|
||||
|
||||
def create_collection(name: str, dim: int) -> Collection:
|
||||
"""创建或获取 collection"""
|
||||
if utility.has_collection(name):
|
||||
print(f"Collection '{name}' 已存在，删除重建")
|
||||
utility.drop_collection(name)
|
||||
|
||||
fields = [
|
||||
FieldSchema(name="chunk_id", dtype=DataType.VARCHAR, max_length=64, is_primary=True),
|
||||
FieldSchema(name="doc_id", dtype=DataType.VARCHAR, max_length=128),
|
||||
FieldSchema(name="doc_title", dtype=DataType.VARCHAR, max_length=512),
|
||||
FieldSchema(name="chunk_index", dtype=DataType.INT64),
|
||||
FieldSchema(name="semantic_id", dtype=DataType.VARCHAR, max_length=64),
|
||||
FieldSchema(name="chunk_type", dtype=DataType.VARCHAR, max_length=32),
|
||||
FieldSchema(name="page_start", dtype=DataType.INT64),
|
||||
FieldSchema(name="page_end", dtype=DataType.INT64),
|
||||
FieldSchema(name="section_title", dtype=DataType.VARCHAR, max_length=512),
|
||||
FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=2048),
|
||||
FieldSchema(name="source_ids", dtype=DataType.VARCHAR, max_length=4096), # JSON 字符串
|
||||
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
||||
]
|
||||
|
||||
schema = CollectionSchema(fields, description="法规文档检索 chunks")
|
||||
collection = Collection(name, schema)
|
||||
|
||||
# 创建向量索引（IVF_FLAT，适合中小规模）
|
||||
index_params = {
|
||||
"metric_type": "COSINE",
|
||||
"index_type": "IVF_FLAT",
|
||||
"params": {"nlist": 128},
|
||||
}
|
||||
collection.create_index("embedding", index_params)
|
||||
print(f"Collection '{name}' 创建完成，索引已建立")
|
||||
|
||||
return collection
|
||||
|
||||
|
||||
def insert_chunks(collection: Collection, chunks: List[Dict], embeddings: List[List[float]]):
|
||||
"""插入 chunks 到 Milvus"""
|
||||
data = [
|
||||
[c["chunk_id"] for c in chunks],
|
||||
[c["doc_id"] for c in chunks],
|
||||
[c["doc_title"] for c in chunks],
|
||||
[c["chunk_index"] for c in chunks],
|
||||
[c["semantic_id"] for c in chunks],
|
||||
[c["chunk_type"] for c in chunks],
|
||||
[c["page_start"] for c in chunks],
|
||||
[c["page_end"] for c in chunks],
|
||||
[c["section_title"] for c in chunks],
|
||||
[c["text"] for c in chunks],
|
||||
[json.dumps(c.get("source_ids", [])) for c in chunks], # JSON 字符串
|
||||
embeddings,
|
||||
]
|
||||
|
||||
collection.insert(data)
|
||||
collection.flush()
|
||||
print(f"已插入 {len(chunks)} 个 chunks")
|
||||
|
||||
|
||||
def load_collection(collection: Collection):
|
||||
"""加载 collection 到内存（搜索前必须）"""
|
||||
collection.load()
|
||||
print(f"Collection 已加载到内存")
|
||||
|
||||
|
||||
# ===================== PostgreSQL =====================
|
||||
def get_pg_connection(host: str, port: int, user: str, password: str, database: str):
|
||||
"""获取 PostgreSQL 连接"""
|
||||
conn = psycopg2.connect(
|
||||
host=host,
|
||||
port=port,
|
||||
user=user,
|
||||
password=password,
|
||||
database=database,
|
||||
)
|
||||
print(f"已连接 PostgreSQL: {host}:{port}/{database}")
|
||||
return conn
|
||||
|
||||
|
||||
def insert_chunks_to_pg(conn, chunks: List[Dict], doc_data: Dict):
|
||||
"""插入 chunks 和相关数据到 PostgreSQL"""
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
# 1. 插入文档
|
||||
cursor.execute("""
|
||||
INSERT INTO documents (doc_id, title, standard_number, upload_time)
|
||||
VALUES (%s, %s, %s, NOW())
|
||||
ON CONFLICT (doc_id) DO UPDATE SET title = EXCLUDED.title, updated_at = NOW()
|
||||
""", (doc_data["doc_id"], doc_data["doc_title"], doc_data.get("standard_number")))
|
||||
|
||||
# 2. 插入语义块
|
||||
semantic_blocks = doc_data.get("semantic_blocks", [])
|
||||
if semantic_blocks:
|
||||
block_rows = [
|
||||
(
|
||||
doc_data["doc_id"],
|
||||
block["semantic_id"],
|
||||
block["block_type"],
|
||||
block["page_start"],
|
||||
block["page_end"],
|
||||
block.get("section_title"),
|
||||
block.get("section_level"),
|
||||
json.dumps(block.get("source_ids", [])),
|
||||
block["text"],
|
||||
)
|
||||
for block in semantic_blocks
|
||||
]
|
||||
execute_values(
|
||||
cursor,
|
||||
"""
|
||||
INSERT INTO semantic_blocks
|
||||
(doc_id, semantic_id, block_type, page_start, page_end, section_title, section_level, source_ids, text)
|
||||
VALUES %s
|
||||
ON CONFLICT (doc_id, semantic_id) DO UPDATE SET text = EXCLUDED.text
|
||||
""",
|
||||
block_rows,
|
||||
)
|
||||
print(f"已插入 {len(semantic_blocks)} 个语义块")
|
||||
|
||||
# 3. 插入向量块元数据
|
||||
chunk_rows = [
|
||||
(
|
||||
doc_data["doc_id"],
|
||||
chunk["chunk_id"],
|
||||
chunk["semantic_id"],
|
||||
chunk["chunk_index"],
|
||||
chunk.get("piece_index"),
|
||||
chunk["page_start"],
|
||||
chunk["page_end"],
|
||||
chunk.get("section_title"),
|
||||
chunk["text"],
|
||||
json.dumps(chunk.get("source_ids", [])),
|
||||
)
|
||||
for chunk in chunks
|
||||
]
|
||||
execute_values(
|
||||
cursor,
|
||||
"""
|
||||
INSERT INTO vector_chunks
|
||||
(doc_id, chunk_id, semantic_id, chunk_index, piece_index, page_start, page_end, section_title, text, source_ids)
|
||||
VALUES %s
|
||||
ON CONFLICT (doc_id, chunk_id) DO UPDATE SET text = EXCLUDED.text
|
||||
""",
|
||||
chunk_rows,
|
||||
)
|
||||
print(f"已插入 {len(chunks)} 个向量块元数据")
|
||||
|
||||
conn.commit()
|
||||
print("PostgreSQL 数据插入完成")
|
||||
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
raise e
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
|
||||
# ===================== 主流程 =====================
|
||||
def load_data(file_path: Path) -> Dict:
|
||||
"""加载 vector_chunks.json，返回完整数据"""
|
||||
data = json.loads(file_path.read_text(encoding="utf-8"))
|
||||
return data
|
||||
|
||||
|
||||
def upload_to_milvus_and_pg(
|
||||
chunks_file: str,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
milvus_host: str,
|
||||
milvus_port: str,
|
||||
collection_name: str,
|
||||
batch_size: int,
|
||||
pg_host: str,
|
||||
pg_port: int,
|
||||
pg_user: str,
|
||||
pg_password: str,
|
||||
pg_database: str,
|
||||
):
|
||||
# 1. 加载完整数据
|
||||
chunks_path = Path(chunks_file).expanduser().resolve()
|
||||
if not chunks_path.exists():
|
||||
raise FileNotFoundError(f"文件不存在: {chunks_path}")
|
||||
|
||||
data = load_data(chunks_path)
|
||||
chunks = data.get("vector_chunks", [])
|
||||
if not chunks:
|
||||
raise ValueError("vector_chunks 为空")
|
||||
print(f"加载 {len(chunks)} 个 chunks")
|
||||
|
||||
# 2. 初始化连接
|
||||
client = get_openai_client(api_key, base_url)
|
||||
init_milvus(milvus_host, milvus_port)
|
||||
pg_conn = get_pg_connection(pg_host, pg_port, pg_user, pg_password, pg_database)
|
||||
|
||||
# 3. 获取 embeddings
|
||||
texts = [c["embedding_text"] for c in chunks]
|
||||
embeddings = get_embeddings_batch(client, texts, batch_size)
|
||||
print(f"生成 {len(embeddings)} 个向量")
|
||||
|
||||
# 4. 获取 embedding 维度
|
||||
embedding_dim = len(embeddings[0])
|
||||
print(f"Embedding 维度: {embedding_dim}")
|
||||
|
||||
# 5. 创建 collection 并插入 Milvus
|
||||
collection = create_collection(collection_name, embedding_dim)
|
||||
insert_chunks(collection, chunks, embeddings)
|
||||
load_collection(collection)
|
||||
|
||||
# 6. 插入 PostgreSQL
|
||||
insert_chunks_to_pg(pg_conn, chunks, data)
|
||||
|
||||
# 7. 关闭连接
|
||||
pg_conn.close()
|
||||
|
||||
print("上传完成！")
|
||||
|
||||
|
||||
# ===================== CLI =====================
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="将 vector_chunks 向量化并上传到 Milvus 和 PostgreSQL")
|
||||
parser.add_argument("chunks_file", help="vector_chunks.json 文件路径")
|
||||
parser.add_argument("--api-key", default=RELAY_API_KEY, help="中转站 API Key")
|
||||
parser.add_argument("--base-url", default=RELAY_BASE_URL, help="中转站 Base URL")
|
||||
parser.add_argument("--milvus-host", default=MILVUS_HOST, help="Milvus host")
|
||||
parser.add_argument("--milvus-port", default=MILVUS_PORT, help="Milvus port")
|
||||
parser.add_argument("--collection", default=COLLECTION_NAME, help="Milvus collection 名称")
|
||||
parser.add_argument("--batch-size", type=int, default=10, help="Embedding 批量大小（中转站限制最大10）")
|
||||
parser.add_argument("--pg-host", default=PG_HOST, help="PostgreSQL host")
|
||||
parser.add_argument("--pg-port", type=int, default=PG_PORT, help="PostgreSQL port")
|
||||
parser.add_argument("--pg-user", default=PG_USER, help="PostgreSQL user")
|
||||
parser.add_argument("--pg-password", default=PG_PASSWORD, help="PostgreSQL password")
|
||||
parser.add_argument("--pg-database", default=PG_DATABASE, help="PostgreSQL database")
|
||||
args = parser.parse_args()
|
||||
|
||||
upload_to_milvus_and_pg(
|
||||
chunks_file=args.chunks_file,
|
||||
api_key=args.api_key,
|
||||
base_url=args.base_url,
|
||||
milvus_host=args.milvus_host,
|
||||
milvus_port=args.milvus_port,
|
||||
collection_name=args.collection,
|
||||
batch_size=args.batch_size,
|
||||
pg_host=args.pg_host,
|
||||
pg_port=args.pg_port,
|
||||
pg_user=args.pg_user,
|
||||
pg_password=args.pg_password,
|
||||
pg_database=args.pg_database,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
5212
backend/aliyun_parser/vector_chunks.json
Normal file
5212
backend/aliyun_parser/vector_chunks.json
Normal file
File diff suppressed because it is too large
Load Diff
263
backend/aliyun_parser/嵌入和召回.md
Normal file
263
backend/aliyun_parser/嵌入和召回.md
Normal file
@@ -0,0 +1,263 @@
|
||||
# 文档解析与向量检索说明
|
||||
|
||||
## 相关文件
|
||||
|
||||
- `aliyun_doc_parser.py`：调用阿里云文档智能解析 PDF，生成原始 `layouts.json`
|
||||
- `layouts_to_vector_chunks.py`：把 `layouts.json` 转成适合向量数据库入库的三层结构
|
||||
- `layouts.json`：阿里云返回的原始布局结果
|
||||
- `vector_chunks.json`：转换后的结构化输出
|
||||
|
||||
## 一、`layouts.json` 的结构
|
||||
|
||||
`layouts.json` 顶层是一个数组，每个元素代表一个布局块（layout）。常见字段如下：
|
||||
|
||||
- `type`：主类型，例如 `title`、`text`、`table`、`figure`
|
||||
- `subType`：更细的语义类型，例如 `doc_title`、`para_title`、`para`、`picture`、`pic_title`、`pic_caption`
|
||||
- `text`：当前布局块的纯文本
|
||||
- `markdownContent`：带 markdown 标记的文本
|
||||
- `pageNum`：页码
|
||||
- `index`：页内顺序
|
||||
- `level`：标题层级
|
||||
- `uniqueId`：布局块唯一标识
|
||||
- `blocks`：更细粒度的文本与样式信息
|
||||
- `cells`：表格单元格，仅 `table` 类型存在
|
||||
|
||||
这个结构不是简单 OCR 文本流，而是已经带有版面理解和语义分类的结构化数据。
|
||||
|
||||
## 二、推荐的三层转换结构
|
||||
|
||||
### 1. 结构层 `structure_nodes`
|
||||
|
||||
结构层用于恢复文档标题树，不直接作为最终向量检索单元。
|
||||
|
||||
示例：
|
||||
|
||||
- `1 范围`
|
||||
- `2 规范性引用文件`
|
||||
- `3 术语和定义`
|
||||
- `3.1 儿童三轮车`
|
||||
- `3.2 轮距`
|
||||
|
||||
结构层主要用于给下游 chunk 绑定 `section_path`。
|
||||
|
||||
### 2. 语义层 `semantic_blocks`
|
||||
|
||||
语义层是按文档意义聚合后的内容块，主要分为三类：
|
||||
|
||||
- `section_text`：同一章节下连续正文聚合而成
|
||||
- `table`：表格内容单独成块
|
||||
- `figure`：图、图名、图注等单独成块
|
||||
|
||||
这一层比单 layout 更适合做语义理解，也适合后续做上下文扩展。
|
||||
|
||||
### 3. 检索层 `vector_chunks`
|
||||
|
||||
检索层是最终写进向量数据库的 chunk。
|
||||
|
||||
处理方式：
|
||||
|
||||
- 对 `semantic_blocks` 中较短的块直接入库
|
||||
- 对较长的块按 `max_chars` 再切分
|
||||
- 相邻切片保留 `overlap_chars` 重叠
|
||||
- 每个 chunk 都带完整 metadata，便于后续过滤、重排和邻域扩展
|
||||
|
||||
## 三、当前转换脚本做了什么
|
||||
|
||||
`layouts_to_vector_chunks.py` 当前已经实现：
|
||||
|
||||
1. 过滤目录页噪声（如 `目次`）
|
||||
2. 根据标题层级维护章节路径
|
||||
3. 将正文聚合成 `section_text`
|
||||
4. 将表格单独转成 `table`
|
||||
5. 将图相关内容单独转成 `figure`
|
||||
6. 对长文本继续切分为最终 `vector_chunks`
|
||||
7. 为每个检索 chunk 生成 `embedding_text`
|
||||
|
||||
## 四、为什么不要直接按 layout 入库
|
||||
|
||||
如果把 `layouts.json` 的每条 layout 直接做向量：
|
||||
|
||||
- 颗粒度太碎
|
||||
- 标题和正文容易分离
|
||||
- 表格会丢失结构上下文
|
||||
- 图示信息无法完整表达
|
||||
- 检索命中结果噪声较大
|
||||
|
||||
对于标准文档，最合适的单位通常不是“句子”，而是“条款语义块”。
|
||||
|
||||
## 五、建议的入库字段
|
||||
|
||||
建议向量数据库每条记录至少保存：
|
||||
|
||||
- `embedding_text`：用于生成向量
|
||||
- `text`：原始 chunk 文本
|
||||
- `chunk_id`
|
||||
- `semantic_id`
|
||||
- `chunk_type`：`section_text` / `table` / `figure`
|
||||
- `section_path`
|
||||
- `section_title`
|
||||
- `section_level`
|
||||
- `page_start`
|
||||
- `page_end`
|
||||
- `doc_id`
|
||||
- `doc_title`
|
||||
- `source_ids`
|
||||
|
||||
其中：
|
||||
|
||||
- 向量化字段：`embedding_text`
|
||||
- 展示字段：`text`
|
||||
- 检索增强字段：其余 metadata
|
||||
|
||||
## 六、推荐的检索方式
|
||||
|
||||
不要只做最简单的 top-k 向量搜索，建议采用：
|
||||
|
||||
**向量召回 + metadata 重排 + 邻域扩展**
|
||||
|
||||
### 1. 向量召回
|
||||
|
||||
使用 `vector_chunks[*].embedding_text` 做 embedding，并在向量数据库中检索 top 10 ~ 15 条。
|
||||
|
||||
查询时可以对用户问题做轻微改写，例如：
|
||||
|
||||
原问题：
|
||||
|
||||
`儿童三轮车的定义是什么？`
|
||||
|
||||
可改写为：
|
||||
|
||||
`请检索 GB 14747—2006 儿童三轮车安全要求 中关于“儿童三轮车定义”的条款、术语、表格或图示说明。`
|
||||
|
||||
这样更适合标准文档检索。
|
||||
|
||||
### 2. metadata 重排
|
||||
|
||||
向量召回后，根据 metadata 做轻量规则重排。
|
||||
|
||||
常见规则：
|
||||
|
||||
- `chunk_type == section_text`：对定义类、要求类问题优先级更高
|
||||
- `section_path` 命中查询关键词：例如查询“定义”时，`术语和定义` 章节优先
|
||||
- `chunk_type == table`：对“尺寸 / 参数 / 数值 / 对照 / 要求”类问题加权
|
||||
- `chunk_type == figure`：对“图 / 结构 / 状态 / 示意”类问题加权
|
||||
|
||||
### 3. 邻域扩展
|
||||
|
||||
检索命中的是最终切片，但回答往往需要更完整上下文。
|
||||
|
||||
建议命中某个 `vector_chunk` 后：
|
||||
|
||||
1. 优先回捞同一个 `semantic_id` 下的所有 chunk
|
||||
2. 如果还不够，再补充同 `section_path`、相邻页码或相邻 `chunk_index` 的内容
|
||||
|
||||
这样可以恢复完整条款，而不是只给模型一小段碎片。
|
||||
|
||||
## 七、不同问题的检索重点
|
||||
|
||||
### 1. 定义类问题
|
||||
|
||||
例如：
|
||||
|
||||
- `儿童三轮车的定义是什么？`
|
||||
- `轮距是什么意思？`
|
||||
|
||||
优先检索：
|
||||
|
||||
- `section_text`
|
||||
- `section_path` 中包含 `术语和定义` 的内容
|
||||
|
||||
### 2. 要求类问题
|
||||
|
||||
例如：
|
||||
|
||||
- `外露突出物有什么要求？`
|
||||
- `辅助推杆有哪些安全要求？`
|
||||
|
||||
优先检索：
|
||||
|
||||
- `section_text`
|
||||
- `table`
|
||||
|
||||
### 3. 数值 / 尺寸 / 对照类问题
|
||||
|
||||
例如：
|
||||
|
||||
- `鞍座到脚蹬距离要求是什么？`
|
||||
- `哪些项目需要满足规定尺寸？`
|
||||
|
||||
优先检索：
|
||||
|
||||
- `table`
|
||||
- `section_text`
|
||||
|
||||
### 4. 图示说明类问题
|
||||
|
||||
例如：
|
||||
|
||||
- `正常乘骑状态是什么意思？`
|
||||
- `图1表示什么？`
|
||||
|
||||
优先检索：
|
||||
|
||||
- `figure`
|
||||
- 同章节相邻 `section_text`
|
||||
|
||||
## 八、推荐的最终检索流程
|
||||
|
||||
建议采用以下固定流程：
|
||||
|
||||
1. 用 `vector_chunks.embedding_text` 做 embedding 检索
|
||||
2. 取 top 10 ~ 15 条候选
|
||||
3. 按 `chunk_type + section_path` 做规则重排
|
||||
4. 以 `semantic_id` 为中心回捞完整语义块
|
||||
5. 选 3 ~ 5 组上下文提供给大模型回答
|
||||
|
||||
## 九、给大模型的上下文组织方式
|
||||
|
||||
最终不要直接把原始 JSON 扔给模型，建议整理成如下格式：
|
||||
|
||||
```text
|
||||
[命中片段 1]
|
||||
章节：3 术语和定义 > 3.1 儿童三轮车
|
||||
页码：1-2
|
||||
类型：section_text
|
||||
内容：
|
||||
......
|
||||
|
||||
[命中片段 2]
|
||||
章节：4 要求 > 4.3 外露突出物
|
||||
页码：5
|
||||
类型：section_text
|
||||
内容：
|
||||
......
|
||||
|
||||
[命中片段 3]
|
||||
章节：5 试验方法
|
||||
页码：8
|
||||
类型：table
|
||||
内容：
|
||||
......
|
||||
```
|
||||
|
||||
这种格式更利于模型稳定回答并引用出处。
|
||||
|
||||
## 十、转换命令
|
||||
|
||||
生成三层结构：
|
||||
|
||||
```bash
|
||||
python3 /home/huaci/dev/ai/SuperMew/tests/layouts_to_vector_chunks.py \
|
||||
--layouts /home/huaci/dev/ai/SuperMew/tests/layouts.json \
|
||||
--out /home/huaci/dev/ai/SuperMew/tests/vector_chunks.json
|
||||
```
|
||||
|
||||
自定义切片大小：
|
||||
|
||||
```bash
|
||||
python3 /home/huaci/dev/ai/SuperMew/tests/layouts_to_vector_chunks.py \
|
||||
--layouts /home/huaci/dev/ai/SuperMew/tests/layouts.json \
|
||||
--out /home/huaci/dev/ai/SuperMew/tests/vector_chunks.json \
|
||||
--max-chars 500 \
|
||||
--overlap-chars 80
|
||||
```
|
||||
Reference in New Issue
Block a user