332 lines
12 KiB
Python
332 lines
12 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
按舰船号拆分 kg_snapshot.json、kg_output/、files/ 等目录
|
||
输出结果放在 OUTPUT_BASE_DIR 下，每艘舰一个子目录
|
||
|
||
文件名规则与舰船的对应关系只需修改下方「变量配置区」，逻辑不用动。
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import re
|
||
import shutil
|
||
import sys
|
||
from collections import defaultdict
|
||
from typing import Dict, List, Optional, Set
|
||
|
||
# 从统一配置文件读取路径
|
||
sys.path.insert(0, os.path.dirname(__file__))
|
||
from migration_config import DST_CODE_DIR, DST_FILES_DIR
|
||
|
||
# ===================== 变量配置区（路径已从 migration_config.py 读取） ========================
|
||
|
||
# 输入路径
|
||
SNAPSHOT_PATH = os.path.join(DST_CODE_DIR, "kg_snapshot.json") # 全量快照
|
||
KG_OUTPUT_DIR = os.path.join(DST_CODE_DIR, "kg_output") # 抽取结果 JSON 目录
|
||
FILES_DIR = DST_FILES_DIR # 原始文档目录
|
||
UPLOADS_V1_DIR = os.path.join(DST_FILES_DIR, "uploads_v1") # 高亮 PDF 目录
|
||
|
||
# 输出根目录（每艘舰一个子目录）
|
||
OUTPUT_BASE_DIR = os.path.join(DST_CODE_DIR, "migration_split")
|
||
|
||
# --------------------------------------------------------------------------
|
||
# 舰船规则表：定义「哪些文件名属于哪艘舰」
|
||
#
|
||
# 每条规则是一个 dict，支持以下匹配方式（可以组合，任意一条命中即匹配）：
|
||
# prefix : 文件名以此开头（不含路径）
|
||
# suffix : 文件名以此结尾（含扩展名）
|
||
# contains : 文件名包含此字符串
|
||
# regex : 正则表达式，匹配整个文件名（不含路径）
|
||
#
|
||
# 示例：
|
||
# {"prefix": "122-"} → 以 "122-" 开头
|
||
# {"prefix": "163-"} → 以 "163-" 开头
|
||
# {"contains": "发动机"} → 包含"发动机"
|
||
# {"regex": r"^JZ\d+_.+"} → 正则匹配
|
||
# --------------------------------------------------------------------------
|
||
SHIP_RULES: Dict[str, List[Dict]] = {
|
||
"101舰": [
|
||
{"prefix": "101-"},
|
||
],
|
||
"122舰": [
|
||
{"prefix": "122-"},
|
||
],
|
||
"163舰": [
|
||
{"prefix": "163-"},
|
||
],
|
||
# 新增舰船示例（取消注释并修改）：
|
||
# "200舰": [
|
||
# {"prefix": "200-"},
|
||
# {"contains": "某特殊前缀"},
|
||
# ],
|
||
}
|
||
|
||
# 共享节点（同时属于多艘舰）的处理策略：
|
||
# "all" : 每艘涉及的舰都保留这个节点（knowledge_source 保持完整）
|
||
# "majority" : 归属来源文件更多的那艘舰，其他舰不保留
|
||
SHARED_NODE_POLICY = "all"
|
||
|
||
# ==========================================================================
|
||
|
||
|
||
def match_ship(filename: str) -> Optional[str]:
|
||
"""
|
||
根据 SHIP_RULES 判断文件名属于哪艘舰。
|
||
返回舰船名，或 None（没有匹配）。
|
||
第一条命中的规则生效（规则之间互斥时注意顺序）。
|
||
"""
|
||
bare = os.path.basename(filename)
|
||
for ship, rules in SHIP_RULES.items():
|
||
for rule in rules:
|
||
if "prefix" in rule and bare.startswith(rule["prefix"]):
|
||
return ship
|
||
if "suffix" in rule and bare.endswith(rule["suffix"]):
|
||
return ship
|
||
if "contains" in rule and rule["contains"] in bare:
|
||
return ship
|
||
if "regex" in rule and re.match(rule["regex"], bare):
|
||
return ship
|
||
return None
|
||
|
||
|
||
# --------------------------------------------------------------------------
|
||
# 快照拆分
|
||
# --------------------------------------------------------------------------
|
||
|
||
def _get_node_ships(entities: List[Dict]) -> Dict[str, Set[str]]:
|
||
"""返回 {节点名: {归属舰集合}}"""
|
||
node_ships: Dict[str, Set[str]] = defaultdict(set)
|
||
for e in entities:
|
||
name = e.get("properties", {}).get("名称", "")
|
||
if not name:
|
||
continue
|
||
ks_raw = e.get("properties", {}).get("knowledge_source", "")
|
||
try:
|
||
ks_list = json.loads(ks_raw) if isinstance(ks_raw, str) else (ks_raw or [])
|
||
except Exception:
|
||
ks_list = []
|
||
for ks in ks_list:
|
||
ship = match_ship(ks.get("filename", ""))
|
||
if ship:
|
||
node_ships[name].add(ship)
|
||
return node_ships
|
||
|
||
|
||
def split_snapshot(snapshot_path: str, output_base: str):
|
||
print(f"\n[快照拆分] 读取 {snapshot_path} ...")
|
||
with open(snapshot_path, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
|
||
entities = data.get("entities", [])
|
||
relationships = data.get("relationships", [])
|
||
print(f" 总节点: {len(entities)} 总关系: {len(relationships)}")
|
||
|
||
node_ships = _get_node_ships(entities)
|
||
|
||
# 统计共享情况
|
||
all_ships = set(SHIP_RULES.keys())
|
||
stats = defaultdict(int)
|
||
for name, ships in node_ships.items():
|
||
if len(ships) == 0:
|
||
stats["无来源"] += 1
|
||
elif len(ships) == 1:
|
||
stats[next(iter(ships))] += 1
|
||
else:
|
||
stats["共享"] += 1
|
||
print(f" 节点归属统计: { dict(stats) }")
|
||
|
||
# 为每艘舰生成节点/关系子集
|
||
for ship in SHIP_RULES:
|
||
ship_entities = []
|
||
for e in entities:
|
||
name = e.get("properties", {}).get("名称", "")
|
||
ships = node_ships.get(name, set())
|
||
|
||
if ship in ships:
|
||
ship_entities.append(e)
|
||
elif not ships and SHARED_NODE_POLICY == "all":
|
||
# 无来源文件的节点（如顶层舰艇/系统节点）每艘舰都保留
|
||
ship_entities.append(e)
|
||
|
||
# 只保留两端节点都在本舰集合里的关系
|
||
ship_node_names = {e.get("properties", {}).get("名称", "") for e in ship_entities}
|
||
ship_rels = [
|
||
r for r in relationships
|
||
if r.get("from_entity") in ship_node_names
|
||
and r.get("to_entity") in ship_node_names
|
||
]
|
||
|
||
out_dir = os.path.join(output_base, ship)
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
out_path = os.path.join(out_dir, "kg_snapshot.json")
|
||
with open(out_path, "w", encoding="utf-8") as f:
|
||
json.dump({"entities": ship_entities, "relationships": ship_rels},
|
||
f, ensure_ascii=False, indent=2)
|
||
print(f" [{ship}] 节点:{len(ship_entities)} 关系:{len(ship_rels)} → {out_path}")
|
||
|
||
|
||
# --------------------------------------------------------------------------
|
||
# kg_output 拆分（每个 .json 文件对应一个原始文档）
|
||
# --------------------------------------------------------------------------
|
||
|
||
def split_kg_output(kg_output_dir: str, output_base: str):
|
||
if not os.path.isdir(kg_output_dir):
|
||
print(f"\n[kg_output 拆分] 目录不存在，跳过: {kg_output_dir}")
|
||
return
|
||
|
||
print(f"\n[kg_output 拆分] 扫描 {kg_output_dir} ...")
|
||
files = [f for f in os.listdir(kg_output_dir) if f.endswith(".json")]
|
||
print(f" 共 {len(files)} 个文件")
|
||
|
||
counts = defaultdict(int)
|
||
unmatched = []
|
||
|
||
for fname in files:
|
||
# kg_output 文件名格式：原始文件名.json（含扩展名，如 122-xxx.pdf.json）
|
||
# 去掉最后的 .json 拿到原始文件名
|
||
original = fname[:-5] if fname.endswith(".json") else fname
|
||
ship = match_ship(original)
|
||
|
||
if ship:
|
||
dst_dir = os.path.join(output_base, ship, "kg_output")
|
||
os.makedirs(dst_dir, exist_ok=True)
|
||
shutil.copy2(os.path.join(kg_output_dir, fname),
|
||
os.path.join(dst_dir, fname))
|
||
counts[ship] += 1
|
||
else:
|
||
unmatched.append(fname)
|
||
|
||
for ship, cnt in counts.items():
|
||
print(f" [{ship}] {cnt} 个 JSON")
|
||
if unmatched:
|
||
print(f" [未匹配] {len(unmatched)} 个文件: {unmatched[:5]}")
|
||
|
||
|
||
# --------------------------------------------------------------------------
|
||
# files/ 原始文档拆分
|
||
# --------------------------------------------------------------------------
|
||
|
||
def split_files(files_dir: str, output_base: str):
|
||
if not os.path.isdir(files_dir):
|
||
print(f"\n[files 拆分] 目录不存在，跳过: {files_dir}")
|
||
return
|
||
|
||
print(f"\n[files 拆分] 扫描 {files_dir} ...")
|
||
all_files = [f for f in os.listdir(files_dir) if os.path.isfile(os.path.join(files_dir, f))]
|
||
print(f" 共 {len(all_files)} 个文件")
|
||
|
||
counts = defaultdict(int)
|
||
unmatched = []
|
||
|
||
for fname in all_files:
|
||
# files/ 里有两种格式：
|
||
# UUID_原文件名.pdf → 去掉 UUID 前缀再匹配
|
||
# 原文件名.pdf → 直接匹配
|
||
bare = fname
|
||
# 尝试去掉 UUID 前缀（格式：8-4-4-4-12 字符的十六进制）
|
||
uuid_pattern = r'^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}_(.+)$'
|
||
m = re.match(uuid_pattern, fname, re.IGNORECASE)
|
||
if m:
|
||
bare = m.group(1)
|
||
|
||
ship = match_ship(bare)
|
||
if ship:
|
||
dst_dir = os.path.join(output_base, ship, "files")
|
||
os.makedirs(dst_dir, exist_ok=True)
|
||
shutil.copy2(os.path.join(files_dir, fname),
|
||
os.path.join(dst_dir, fname))
|
||
counts[ship] += 1
|
||
else:
|
||
unmatched.append(fname)
|
||
|
||
for ship, cnt in counts.items():
|
||
print(f" [{ship}] {cnt} 个文件")
|
||
if unmatched:
|
||
print(f" [未匹配] {len(unmatched)} 个: {unmatched[:5]}")
|
||
|
||
|
||
# --------------------------------------------------------------------------
|
||
# uploads_v1/ 高亮 PDF 拆分（通过快照里的 knowledge_source 反查归属）
|
||
# --------------------------------------------------------------------------
|
||
|
||
def split_uploads(snapshot_path: str, uploads_dir: str, output_base: str):
|
||
if not os.path.isdir(uploads_dir):
|
||
print(f"\n[uploads 拆分] 目录不存在，跳过: {uploads_dir}")
|
||
return
|
||
|
||
print(f"\n[uploads 拆分] 通过快照反查高亮 PDF 归属 ...")
|
||
|
||
with open(snapshot_path, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
|
||
# 从快照里收集：高亮PDF文件名 → 归属舰船
|
||
highlight_ship: Dict[str, str] = {}
|
||
for e in data.get("entities", []):
|
||
ks_raw = e.get("properties", {}).get("knowledge_source", "")
|
||
try:
|
||
ks_list = json.loads(ks_raw) if isinstance(ks_raw, str) else (ks_raw or [])
|
||
except Exception:
|
||
continue
|
||
for ks in ks_list:
|
||
url = ks.get("url", "")
|
||
src_file = ks.get("filename", "")
|
||
ship = match_ship(src_file)
|
||
if not ship or not url:
|
||
continue
|
||
# URL 格式：http://host/upload/<highlight_filename>#page=N
|
||
if "/upload/" in url:
|
||
highlight_fname = url.split("/upload/")[-1].split("#")[0]
|
||
# 共享节点下同一个高亮文件只记录第一个匹配舰船
|
||
if highlight_fname not in highlight_ship:
|
||
highlight_ship[highlight_fname] = ship
|
||
|
||
counts = defaultdict(int)
|
||
missing = []
|
||
|
||
for fname, ship in highlight_ship.items():
|
||
src = os.path.join(uploads_dir, fname)
|
||
if os.path.exists(src):
|
||
dst_dir = os.path.join(output_base, ship, "files", "uploads_v1")
|
||
os.makedirs(dst_dir, exist_ok=True)
|
||
shutil.copy2(src, os.path.join(dst_dir, fname))
|
||
counts[ship] += 1
|
||
else:
|
||
missing.append(fname)
|
||
|
||
for ship, cnt in counts.items():
|
||
print(f" [{ship}] {cnt} 个高亮 PDF")
|
||
if missing:
|
||
print(f" [本地缺失] {len(missing)} 个（需从服务器 uploads_v1 补全）: {missing[:3]}")
|
||
|
||
|
||
# --------------------------------------------------------------------------
|
||
# 入口
|
||
# --------------------------------------------------------------------------
|
||
|
||
def main():
|
||
print("=" * 50)
|
||
print(" kgrag 按舰船号拆分工具")
|
||
print(f" 舰船规则: { {k: v for k, v in SHIP_RULES.items()} }")
|
||
print(f" 共享节点策略: {SHARED_NODE_POLICY}")
|
||
print("=" * 50)
|
||
|
||
os.makedirs(OUTPUT_BASE_DIR, exist_ok=True)
|
||
|
||
split_snapshot(SNAPSHOT_PATH, OUTPUT_BASE_DIR)
|
||
split_kg_output(KG_OUTPUT_DIR, OUTPUT_BASE_DIR)
|
||
split_files(FILES_DIR, OUTPUT_BASE_DIR)
|
||
split_uploads(SNAPSHOT_PATH, UPLOADS_V1_DIR, OUTPUT_BASE_DIR)
|
||
|
||
print("\n" + "=" * 50)
|
||
print(f" 完成！输出目录: {OUTPUT_BASE_DIR}")
|
||
print(" 每艘舰的子目录结构：")
|
||
print(" <舰名>/kg_snapshot.json ← 导入新 Neo4j 用")
|
||
print(" <舰名>/kg_output/ ← 抽取结果 JSON")
|
||
print(" <舰名>/files/ ← 原始文档")
|
||
print(" <舰名>/files/uploads_v1/ ← 高亮 PDF")
|
||
print("=" * 50)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|