423 lines
19 KiB
Python
423 lines
19 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 --folder /app/files --output_dir /app/kg_output --phase all --workers 4
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
import logging
|
||
import argparse
|
||
import asyncio
|
||
import traceback
|
||
import time
|
||
from datetime import datetime
|
||
from multiprocessing import Process, Queue
|
||
from typing import List, Dict, Tuple, Any, Optional
|
||
|
||
EXTRACT_TIMEOUT = 600 # 单文件抽取超时秒数
|
||
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}")
|
||
|
||
# ── 项目内模块 ──────────────────────────────────────────────────────────────
|
||
from extract_text_node_relation import process_file_content, check_filename_ocr
|
||
from kg_build.extract_filtertext import get_filtertext_node_relation
|
||
from filename_proceess_and_kgquery import get_entity
|
||
from kg_build.filter_node_relation import filter_node_relationship
|
||
from kg_build.kg_save import MilitaryKnowledgeGraph
|
||
from modelsAPI.model_api import OpenaiAPI
|
||
import importlib
|
||
import entity_registry as _entity_registry
|
||
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,
|
||
)
|
||
|
||
|
||
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_ANALYZE = os.getenv("MINERU_ANALYZE", "http://192.168.1.64:18000/analyze-pdf")
|
||
MINERU_CACHE = os.getenv("MINERU_CACHE", "http://192.168.1.64:18000/get_content_list")
|
||
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) -> Dict:
|
||
"""
|
||
对单个文件做实体/关系抽取，结果保存到 output_dir/<filename>.json。
|
||
已存在 JSON 则跳过（断点续跑）。
|
||
返回 {"file": ..., "status": "ok"/"skip"/"error", "entity_count": ..., "rel_count": ...}
|
||
"""
|
||
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 = MINERU_ANALYZE,
|
||
check_cancelled_callback = lambda: False,
|
||
device = lower_entity_device,
|
||
device_type = lower_entity_type,
|
||
get_contentlist_url = MINERU_CACHE,
|
||
)
|
||
|
||
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 = [], []
|
||
|
||
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):
|
||
"""子进程入口：抽取单文件，结果放入队列"""
|
||
try:
|
||
result = extract_single_file(file_path, output_dir)
|
||
result_queue.put(result)
|
||
except Exception as e:
|
||
result_queue.put({"file": os.path.basename(file_path), "status": "error", "error": str(e)})
|
||
|
||
|
||
def phase_extract(folder: str, output_dir: str, workers: int):
|
||
"""阶段1：多进程抽取，每个文件独立进程，超时后真正 kill"""
|
||
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}")
|
||
|
||
results = []
|
||
# 按 workers 大小分批，每批同时跑 workers 个进程
|
||
for i in range(0, len(files), workers):
|
||
batch = files[i:i + workers]
|
||
procs = []
|
||
for fp in batch:
|
||
q = Queue()
|
||
p = Process(target=_extract_worker, args=(fp, output_dir, q))
|
||
p.start()
|
||
procs.append((p, fp, q))
|
||
|
||
for p, fp, q in procs:
|
||
file_name = os.path.basename(fp)
|
||
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)
|
||
results.append({"file": file_name, "status": "timeout"})
|
||
else:
|
||
try:
|
||
results.append(q.get_nowait())
|
||
except Exception:
|
||
results.append({"file": file_name, "status": "error"})
|
||
|
||
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("生成 Embedding...")
|
||
t0 = time.time()
|
||
embedded = 0
|
||
for entity in filter_data["entities"]:
|
||
name = entity.get("properties", {}).get("名称", "")
|
||
if not name or "未知" in name:
|
||
continue
|
||
try:
|
||
entity["properties"]["embedding"] = OpenaiAPI.generate_embedding(name)
|
||
embedded += 1
|
||
except Exception as e:
|
||
logger.warning(f"Embedding 失败 [{name}]: {e}")
|
||
logger.info(f"Embedding 完成 {embedded} 个实体 耗时:{time.time()-t0:.1f}s")
|
||
|
||
# ── 4. 写入 Neo4j ─────────────────────────────────────────────────────────
|
||
logger.info("写入 Neo4j...")
|
||
t0 = time.time()
|
||
kg = MilitaryKnowledgeGraph(
|
||
uri = NEO4J_URI,
|
||
auth = (NEO4J_USER, NEO4J_PASSWORD),
|
||
Ontology = 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 main():
|
||
parser = argparse.ArgumentParser(description="批量知识图谱构建")
|
||
parser.add_argument("--folder", default="", help="待处理文件夹（阶段1必填）")
|
||
parser.add_argument("--output_dir", default="/app/kg_output", help="JSON 输出/读取目录")
|
||
parser.add_argument("--phase", default="all", choices=["extract", "merge", "all"],
|
||
help="extract=只抽取 merge=只入库 all=全流程")
|
||
parser.add_argument("--workers", type=int, default=1, help="阶段1并发线程数")
|
||
args = parser.parse_args()
|
||
|
||
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)
|
||
|
||
|
||
def _fake_worker(file_path: str, output_dir: str, result_queue: Queue):
|
||
"""测试用子进程：file_B 睡 20s 模拟超时，其余 2s"""
|
||
file_name = os.path.basename(file_path)
|
||
sleep_time = 20 if file_name == "file_B" else 2
|
||
time.sleep(sleep_time)
|
||
out_path = os.path.join(output_dir, file_name + ".json")
|
||
with open(out_path, "w", encoding="utf-8") as f:
|
||
json.dump({"file": file_name}, f)
|
||
result_queue.put({"file": file_name, "status": "ok"})
|
||
|
||
|
||
def _test_timeout():
|
||
"""
|
||
模拟超时跳过逻辑，验证完删掉此函数及 _fake_worker、下面的调用。
|
||
file_A 正常(2s)，file_B 超时(20s)，file_C 正常(2s)
|
||
期望总耗时约 TIMEOUT+2s，file_B 不写 JSON
|
||
"""
|
||
import tempfile
|
||
|
||
global EXTRACT_TIMEOUT
|
||
EXTRACT_TIMEOUT = 5 # 测试用 5 秒
|
||
|
||
output_dir = tempfile.mkdtemp()
|
||
logger.info(f"测试输出目录: {output_dir}")
|
||
|
||
files = ["file_A", "file_B", "file_C"]
|
||
results = []
|
||
t_start = time.time()
|
||
|
||
for i in range(0, len(files), 2):
|
||
batch = files[i:i + 2]
|
||
procs = []
|
||
for f in batch:
|
||
q = Queue()
|
||
p = Process(target=_fake_worker, args=(f, output_dir, q))
|
||
p.start()
|
||
procs.append((p, f, q))
|
||
|
||
for p, f, q in procs:
|
||
p.join(timeout=EXTRACT_TIMEOUT)
|
||
if p.is_alive():
|
||
p.kill()
|
||
p.join()
|
||
logger.warning(f"[TIMEOUT] {f} 超时，进程已 kill，不写 JSON")
|
||
_record_fault(f)
|
||
results.append({"file": f, "status": "timeout"})
|
||
else:
|
||
results.append(q.get_nowait())
|
||
|
||
logger.info(f"总耗时: {time.time()-t_start:.1f}s")
|
||
logger.info(f"结果: {results}")
|
||
logger.info(f"output_dir 中的文件: {os.listdir(output_dir)}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
_test_timeout()
|
||
# main()
|
||
|
||
|