466 lines
21 KiB
Python
466 lines
21 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
批量知识图谱构建脚本 —— 三阶段 ETL
|
||
|
||
阶段1 (extract): 对文件夹内所有文件做实体抽取，每个文件结果存为 JSON
|
||
阶段2 (merge): 汇总所有 JSON，去重，生成 embedding，写入 Neo4j
|
||
|
||
用法:
|
||
# 只跑抽取（4个容器各跑各的文件夹）
|
||
python batch_kg_build.py --folder /app/files --output_dir /app/kg_output --phase extract --workers 4
|
||
|
||
# 只跑入库（所有容器跑完后，跑一次）
|
||
python batch_kg_build.py --output_dir /app/kg_output --phase merge
|
||
|
||
#抽取好的实体关系直接进行存储
|
||
python batch_kg_build.py --phase load_snapshot --snapshot_path /path/to/my_custom_snapshot.json
|
||
|
||
# 全流程一次跑完
|
||
python batch_kg_build.py --folder /app/files --output_dir /app/kg_output --phase all --workers 4
|
||
"""
|
||
|
||
from default_ontology_config import (
|
||
ONTOLOGY as _BASE_ONTOLOGY,
|
||
ONTOLOGY_DESCRIPTION as _BASE_ONTOLOGY_DESC,
|
||
NODE_TYPES as _BASE_NODE_TYPES,
|
||
RELATIONSHIPS as RELATIONSHIPS,
|
||
RELATIONSHIP_TYPES as RELATIONSHIP_TYPES,
|
||
)
|
||
import entity_registry as _entity_registry
|
||
import importlib
|
||
from modelsAPI.model_api import OpenaiAPI
|
||
from kg_build.kg_save import MilitaryKnowledgeGraph
|
||
from kg_build.filter_node_relation import filter_node_relationship
|
||
from filename_proceess_and_kgquery import get_entity
|
||
from kg_build.extract_filtertext import get_filtertext_node_relation
|
||
from extract_text_node_relation import process_file_content, check_filename_ocr
|
||
import os
|
||
import json
|
||
import logging
|
||
import argparse
|
||
import asyncio
|
||
import traceback
|
||
import time
|
||
from datetime import datetime
|
||
from multiprocessing import Process, Queue
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from typing import List, Dict, Tuple, Any, Optional
|
||
|
||
EXTRACT_TIMEOUT = 600 # 单文件抽取超时秒数
|
||
EMBED_WORKERS = 8 # Embedding 并发线程数
|
||
FAULT_DIR = "/app/fault_wenjian"
|
||
|
||
|
||
def _record_fault(file_name: str):
|
||
"""将超时文件名追加到 /app/fault_wenjian/timeout_YYYYMMDD.txt"""
|
||
os.makedirs(FAULT_DIR, exist_ok=True)
|
||
date_str = datetime.now().strftime("%Y%m%d")
|
||
fault_path = os.path.join(FAULT_DIR, f"timeout_{date_str}.txt")
|
||
with open(fault_path, "a", encoding="utf-8") as f:
|
||
f.write(file_name + "\n")
|
||
logger.warning(f"[FAULT] 已记录超时文件: {fault_path}")
|
||
|
||
|
||
# ── 项目内模块 ──────────────────────────────────────────────────────────────
|
||
|
||
|
||
def _build_ontology_config():
|
||
"""
|
||
与 app.py build_kg_config 保持一致：
|
||
default_ontology_config 的静态配置 + entity_registry.json 动态注册的实体类型合并。
|
||
"""
|
||
importlib.reload(_entity_registry)
|
||
reg_list = _entity_registry.get_ontology_list()
|
||
reg_desc = _entity_registry.get_ontology_description()
|
||
reg_nodes = _entity_registry.get_node_types()
|
||
|
||
# 合并：静态配置优先（手写配置不被覆盖）
|
||
ontology = list(dict.fromkeys(_BASE_ONTOLOGY + reg_list)) # 去重保序
|
||
ontology_desc = {**reg_desc, **_BASE_ONTOLOGY_DESC}
|
||
node_types = {**reg_nodes, **_BASE_NODE_TYPES}
|
||
return ontology, ontology_desc, node_types
|
||
|
||
|
||
# 模块加载时构建一次（抽取阶段开始前已有 registry 内容）
|
||
ONTOLOGY, ONTOLOGY_DESC, NODE_TYPES = _build_ontology_config()
|
||
|
||
# ── 配置（与 app.py 保持一致，按需修改） ────────────────────────────────────
|
||
NEO4J_URI = os.getenv("NEO4J_URI", "bolt://localhost:7687")
|
||
NEO4J_USER = os.getenv("NEO4J_USER", "neo4j")
|
||
NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD", "zdht123@")
|
||
MINERU_INSTANCES = [
|
||
# ("http://192.168.0.64:18000/analyze-pdf"),
|
||
#("http://192.168.0.111:9977/analyze-pdf/"),
|
||
# ("http://192.168.0.111:9979/analyze-pdf"),
|
||
("http://192.168.1.64:18000/analyze-pdf/"),
|
||
]
|
||
PREFIX_URL = os.getenv("PREFIX_URL", "http://192.168.1.64:9085")
|
||
|
||
SUPPORTED_EXT = {".pdf", ".md", ".docx", ".xlsx"}
|
||
|
||
# ── 日志 ────────────────────────────────────────────────────────────────────
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s [%(levelname)s] %(message)s",
|
||
handlers=[
|
||
logging.StreamHandler(),
|
||
logging.FileHandler("batch_kg_build.log", encoding="utf-8"),
|
||
],
|
||
)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
# ════════════════════════════════════════════════════════════════════════════
|
||
# 阶段 1：单文件抽取
|
||
# ════════════════════════════════════════════════════════════════════════════
|
||
|
||
|
||
def extract_single_file(file_path: str, output_dir: str, api_url: str = MINERU_INSTANCES[0]) -> Dict:
|
||
"""
|
||
对单个文件做实体/关系抽取，结果保存到 output_dir/<filename>.json。
|
||
已存在 JSON 则跳过（断点续跑）。
|
||
返回 {"file": ..., "status": "ok"/"skip"/"error", "entity_count": ..., "rel_count": ...}
|
||
"""
|
||
logger.info(f'[api_url] {api_url}')
|
||
file_name = os.path.basename(file_path)
|
||
ext = os.path.splitext(file_name)[1].lower()
|
||
out_path = os.path.join(output_dir, file_name + ".json")
|
||
|
||
# ── 断点续跑：已有结果直接跳过 ──────────────────────────────────────────
|
||
if os.path.exists(out_path):
|
||
logger.info(f"[SKIP] 已有结果，跳过: {file_name}")
|
||
return {"file": file_name, "status": "skip"}
|
||
|
||
if ext not in SUPPORTED_EXT:
|
||
logger.warning(f"[SKIP] 不支持的文件类型: {file_name}")
|
||
return {"file": file_name, "status": "skip"}
|
||
|
||
t0 = time.time()
|
||
entities: List[Dict] = []
|
||
relations: List[Dict] = []
|
||
lower_entity_device: Optional[str] = None
|
||
lower_entity_type: Optional[str] = None
|
||
|
||
try:
|
||
# ── 1. 文件名实体提取 ────────────────────────────────────────────────
|
||
try:
|
||
fn_result = get_entity(file_name)
|
||
fn_entities = fn_result.get("entities", [])
|
||
fn_relations = fn_result.get("relationships", [])
|
||
lower_entity = fn_result.get("low_level")
|
||
for e in fn_entities:
|
||
if e.get("properties", {}).get("名称") == lower_entity:
|
||
lower_entity_device = e["properties"].get("名称")
|
||
lower_entity_type = e.get("type")
|
||
break
|
||
except Exception as e:
|
||
logger.warning(f"[{file_name}] 文件名实体提取失败: {e}")
|
||
fn_entities, fn_relations = [], []
|
||
|
||
# ── 2. 正文实体抽取 ──────────────────────────────────────────────────
|
||
check = check_filename_ocr(file_name)
|
||
|
||
if check == "ocr_true":
|
||
entities, relations = process_file_content(
|
||
file_name=file_name,
|
||
file_path=file_path,
|
||
task_id="batch",
|
||
data_dir=os.path.dirname(file_path),
|
||
ontology=ONTOLOGY,
|
||
ontology_desc=ONTOLOGY_DESC,
|
||
relationships=RELATIONSHIPS,
|
||
node_types=NODE_TYPES,
|
||
relationship_types=RELATIONSHIP_TYPES,
|
||
prefix_url=PREFIX_URL,
|
||
api_url=api_url,
|
||
check_cancelled_callback=lambda: False,
|
||
device=lower_entity_device,
|
||
device_type=lower_entity_type,
|
||
)
|
||
|
||
elif check == "ocr_false":
|
||
async def _run_filtertext():
|
||
return await get_filtertext_node_relation(
|
||
input_pdf_path=file_path,
|
||
data=[],
|
||
prefix_url=PREFIX_URL,
|
||
device=lower_entity_device,
|
||
)
|
||
|
||
result = asyncio.run(_run_filtertext())
|
||
if isinstance(result, tuple) and len(result) >= 2:
|
||
entities, relations = result[0], result[1]
|
||
else:
|
||
entities, relations = [], []
|
||
logger.info("无需进行ocr解析，直接构建实体")
|
||
else:
|
||
logger.warning(f"[{file_name}] check_filename_ocr 返回 None，跳过正文抽取")
|
||
|
||
# ── 3. 合并文件名实体 ────────────────────────────────────────────────
|
||
all_entities = entities + fn_entities
|
||
all_relations = relations + fn_relations
|
||
# 过滤掉没有名称的实体
|
||
all_entities = [e for e in all_entities if e.get("properties", {}).get("名称")]
|
||
|
||
# ── 4. 保存 JSON ─────────────────────────────────────────────────────
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
payload = {"entities": all_entities, "relationships": all_relations}
|
||
with open(out_path, "w", encoding="utf-8") as f:
|
||
json.dump(payload, f, ensure_ascii=False, indent=2)
|
||
|
||
cost = time.time() - t0
|
||
logger.info(f"[OK] {file_name} 实体:{len(all_entities)} 关系:{len(all_relations)} 耗时:{cost:.1f}s")
|
||
return {
|
||
"file": file_name,
|
||
"status": "ok",
|
||
"entity_count": len(all_entities),
|
||
"rel_count": len(all_relations),
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.error(f"[ERROR] {file_name}: {e}\n{traceback.format_exc()}")
|
||
return {"file": file_name, "status": "error", "error": str(e)}
|
||
|
||
|
||
def _extract_worker(file_path: str, output_dir: str, result_queue: Queue, api_url: str = MINERU_INSTANCES[0]):
|
||
"""子进程入口：抽取单文件，结果放入队列"""
|
||
try:
|
||
result = extract_single_file(file_path, output_dir, api_url)
|
||
result_queue.put(result)
|
||
except Exception as e:
|
||
result_queue.put({"file": os.path.basename(file_path), "status": "error", "error": str(e)})
|
||
|
||
|
||
def _run_file_with_timeout(args: Tuple) -> Dict:
|
||
"""在线程中启动子进程并监控超时，实现滚动窗口并发。"""
|
||
fp, output_dir, inst = args
|
||
file_name = os.path.basename(fp)
|
||
q = Queue()
|
||
p = Process(target=_extract_worker, args=(fp, output_dir, q, inst))
|
||
p.start()
|
||
p.join(timeout=EXTRACT_TIMEOUT)
|
||
if p.is_alive():
|
||
p.kill()
|
||
p.join()
|
||
logger.warning(f"[TIMEOUT] {file_name} 超过 {EXTRACT_TIMEOUT}s，已跳过")
|
||
_record_fault(file_name)
|
||
return {"file": file_name, "status": "timeout"}
|
||
try:
|
||
return q.get_nowait()
|
||
except Exception:
|
||
return {"file": file_name, "status": "error"}
|
||
|
||
|
||
def phase_extract(folder: str, output_dir: str, workers: int):
|
||
"""阶段1：滚动窗口多进程抽取，每完成一个立即补充下一个。"""
|
||
files = [
|
||
os.path.join(folder, f)
|
||
for f in os.listdir(folder)
|
||
if os.path.isfile(os.path.join(folder, f))
|
||
and os.path.splitext(f)[1].lower() in SUPPORTED_EXT
|
||
]
|
||
logger.info(f"发现 {len(files)} 个文件，并发数: {workers}")
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
tasks = [
|
||
(fp, output_dir, MINERU_INSTANCES[i % len(MINERU_INSTANCES)])
|
||
for i, fp in enumerate(files)
|
||
]
|
||
|
||
with ThreadPoolExecutor(max_workers=workers) as executor:
|
||
results = list(executor.map(_run_file_with_timeout, tasks))
|
||
|
||
ok = sum(1 for r in results if r["status"] == "ok")
|
||
skip = sum(1 for r in results if r["status"] == "skip")
|
||
error = sum(1 for r in results if r["status"] == "error")
|
||
timeout = sum(1 for r in results if r["status"] == "timeout")
|
||
logger.info(f"抽取完成 成功:{ok} 跳过:{skip} 失败:{error} 超时:{timeout}")
|
||
return results
|
||
|
||
|
||
# ════════════════════════════════════════════════════════════════════════════
|
||
# 阶段 2：汇总 → 去重 → Embedding → 写 Neo4j
|
||
# ════════════════════════════════════════════════════════════════════════════
|
||
|
||
|
||
def phase_merge(output_dir: str):
|
||
"""阶段2：读所有 JSON → 去重 → 向量化 → 写 Neo4j"""
|
||
|
||
# ── 1. 读取所有抽取结果 ──────────────────────────────────────────────────
|
||
all_entities: List[Dict] = []
|
||
all_relations: List[Dict] = []
|
||
|
||
json_files = [f for f in os.listdir(output_dir) if f.endswith(".json")]
|
||
logger.info(f"读取 {len(json_files)} 个抽取结果文件")
|
||
|
||
for fname in json_files:
|
||
fpath = os.path.join(output_dir, fname)
|
||
try:
|
||
with open(fpath, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
all_entities.extend(data.get("entities", []))
|
||
all_relations.extend(data.get("relationships", []))
|
||
except Exception as e:
|
||
logger.warning(f"读取失败 {fname}: {e}")
|
||
|
||
logger.info(f"汇总: 实体 {len(all_entities)} 关系 {len(all_relations)}")
|
||
|
||
# ── 2. 去重（filter_node_relationship 只跑这一次） ───────────────────────
|
||
merged_input = {"entities": all_entities, "relationships": all_relations}
|
||
logger.info("开始全局去重（filter_node_relationship）...")
|
||
t0 = time.time()
|
||
try:
|
||
filter_data = filter_node_relationship(
|
||
URI=NEO4J_URI,
|
||
USERNAME=NEO4J_USER,
|
||
PASSWORD=NEO4J_PASSWORD,
|
||
merged_output_data=merged_input,
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"去重失败，使用原始数据: {e}")
|
||
filter_data = merged_input
|
||
logger.info(f"去重完成 实体:{len(filter_data['entities'])} 关系:{len(filter_data['relationships'])} 耗时:{time.time() - t0:.1f}s")
|
||
|
||
# ── 3. 生成 Embedding（并行） ────────────────────────────────────────────
|
||
logger.info(f"生成 Embedding（并发 {EMBED_WORKERS} 线程）...")
|
||
t0 = time.time()
|
||
|
||
def _embed_one(entity: Dict) -> int:
|
||
name = str(entity.get("properties", {}).get("名称", "") or "")
|
||
if not name or "未知" in name:
|
||
return 0
|
||
try:
|
||
entity["properties"]["embedding"] = OpenaiAPI.generate_embedding(name)
|
||
return 1
|
||
except Exception as e:
|
||
logger.warning(f"Embedding 失败 [{name}]: {e}")
|
||
return 0
|
||
|
||
with ThreadPoolExecutor(max_workers=EMBED_WORKERS) as executor:
|
||
embedded = sum(executor.map(_embed_one, filter_data["entities"]))
|
||
logger.info(f"Embedding 完成 {embedded} 个实体 耗时:{time.time() - t0:.1f}s")
|
||
|
||
# ── 3.5 本地保存带 Embedding 的实体关系 ──────────────────────────────────
|
||
snapshot_path = os.path.join(os.path.dirname(output_dir) if output_dir != "/" else "/app", "kg_snapshot.json")
|
||
try:
|
||
with open(snapshot_path, "w", encoding="utf-8") as f:
|
||
json.dump(filter_data, f, ensure_ascii=False, indent=2)
|
||
logger.info(f"本地快照已保存: {snapshot_path} 实体:{len(filter_data['entities'])} 关系:{len(filter_data['relationships'])}")
|
||
except Exception as e:
|
||
logger.warning(f"本地快照保存失败: {e}")
|
||
|
||
# ── 4. 写入 Neo4j ─────────────────────────────────────────────────────────
|
||
logger.info("写入 Neo4j...")
|
||
# ontology, ontology_desc, node_types = _build_ontology_config()
|
||
|
||
# 从实体数据中提取动态类型，追加到静态 ontology 末尾（静态配置优先，数据类型补充）
|
||
data_types = list(dict.fromkeys(
|
||
e["type"] for e in filter_data["entities"] if e.get("type")
|
||
))
|
||
merged_ontology = list(dict.fromkeys( data_types))
|
||
# new_types = [t for t in data_types if t not in set(ontology)]
|
||
# if new_types:
|
||
# logger.info(f"从实体数据补充新类型 {len(new_types)} 个: {new_types}")
|
||
|
||
t0 = time.time()
|
||
kg = MilitaryKnowledgeGraph(
|
||
uri=NEO4J_URI,
|
||
auth=(NEO4J_USER, NEO4J_PASSWORD),
|
||
Ontology=merged_ontology,
|
||
)
|
||
try:
|
||
entity_count, rel_count = kg.process_data(filter_data)
|
||
logger.info(f"Neo4j 写入完成 实体:{entity_count} 关系:{rel_count} 耗时:{time.time() - t0:.1f}s")
|
||
finally:
|
||
kg.close()
|
||
|
||
|
||
def phase_load_snapshot(snapshot_path: str):
|
||
"""
|
||
Directly read kg_snapshot.json and store entities/relations into Neo4j.
|
||
Skips extraction and embedding generation steps.
|
||
"""
|
||
if not os.path.exists(snapshot_path):
|
||
logger.error(f"Snapshot file not found: {snapshot_path}")
|
||
return
|
||
|
||
logger.info(f"Loading snapshot from: {snapshot_path}")
|
||
t0 = time.time()
|
||
|
||
try:
|
||
with open(snapshot_path, "r", encoding="utf-8") as f:
|
||
filter_data = json.load(f)
|
||
|
||
entities = filter_data.get("entities", [])
|
||
relations = filter_data.get("relationships", [])
|
||
logger.info(f"Loaded snapshot: {len(entities)} entities, {len(relations)} relations")
|
||
|
||
if not entities and not relations:
|
||
logger.warning("Snapshot contains no data to load.")
|
||
return
|
||
|
||
# Extract dynamic types from entities to ensure Ontology coverage
|
||
data_types = list(dict.fromkeys(
|
||
e["type"] for e in entities if e.get("type")
|
||
))
|
||
merged_ontology = list(dict.fromkeys(data_types))
|
||
|
||
logger.info(f"Initializing MilitaryKnowledgeGraph with ontology types: {merged_ontology[:5]}...") # Log first few for brevity
|
||
|
||
kg = MilitaryKnowledgeGraph(
|
||
uri=NEO4J_URI,
|
||
auth=(NEO4J_USER, NEO4J_PASSWORD),
|
||
Ontology=merged_ontology,
|
||
)
|
||
|
||
try:
|
||
entity_count, rel_count = kg.process_data(filter_data)
|
||
elapsed = time.time() - t0
|
||
logger.info(f"Neo4j write completed from snapshot. Entities: {entity_count}, Relations: {rel_count}. Time: {elapsed:.1f}s")
|
||
finally:
|
||
kg.close()
|
||
|
||
except json.JSONDecodeError:
|
||
logger.error(f"Failed to decode JSON from {snapshot_path}")
|
||
except Exception as e:
|
||
logger.error(f"Error loading snapshot: {e}", exc_info=True)
|
||
|
||
# ════════════════════════════════════════════════════════════════════════════
|
||
# 入口
|
||
# ════════════════════════════════════════════════════════════════════════════
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="批量知识图谱构建")
|
||
parser.add_argument("--folder", default="", help="待处理文件夹（阶段1必填）")
|
||
parser.add_argument("--output_dir", default="/app/kg_output", help="JSON 输出/读取目录")
|
||
# Update choices to include 'load_snapshot'
|
||
parser.add_argument("--phase", default="all", choices=["extract", "merge", "load_snapshot", "all"],
|
||
help="extract=只抽取 merge=只入库(含抽取结果汇总) load_snapshot=直接加载快照 all=全流程")
|
||
parser.add_argument("--workers", type=int, default=1, help="阶段1并发线程数")
|
||
parser.add_argument("--snapshot_path", default="", help="快照文件路径 (用于 load_snapshot 阶段)")
|
||
args = parser.parse_args()
|
||
|
||
# Determine default snapshot path if not provided and phase is load_snapshot
|
||
if args.phase == "load_snapshot" and not args.snapshot_path:
|
||
# Default to the location where phase_merge saves it
|
||
args.snapshot_path = os.path.join(os.path.dirname(args.output_dir) if args.output_dir != "/" else "/app", "kg_snapshot.json")
|
||
|
||
if args.phase in ("extract", "all"):
|
||
if not args.folder:
|
||
parser.error("--phase extract/all 需要指定 --folder")
|
||
phase_extract(args.folder, args.output_dir, args.workers)
|
||
|
||
if args.phase in ("merge", "all"):
|
||
phase_merge(args.output_dir)
|
||
|
||
if args.phase == "load_snapshot":
|
||
phase_load_snapshot(args.snapshot_path)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
# ontology, ontology_desc, node_types = _build_ontology_config()
|
||
# print(ontology)
|
||
|