127 lines
3.4 KiB
Python
127 lines
3.4 KiB
Python
#!/usr/bin/env python
|
|
"""
|
|
知识图谱构建脚本
|
|
示例用法：python scripts/build_kg.py --data-dir ../data/1法律 --output ./output/kg.json
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# 添加src到路径
|
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
|
|
|
from src.preprocessing.document_parser import DocumentParser
|
|
from src.kg_builder.indexer import GraphIndexer
|
|
from src.kg_builder.graph import KnowledgeGraph
|
|
from src.utils.config import Config, load_config
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="构建法规知识图谱")
|
|
parser.add_argument(
|
|
"--data-dir",
|
|
type=str,
|
|
required=True,
|
|
help="法规数据目录路径"
|
|
)
|
|
parser.add_argument(
|
|
"--output",
|
|
type=str,
|
|
default="output/kg.json",
|
|
help="输出文件路径"
|
|
)
|
|
parser.add_argument(
|
|
"--max-docs",
|
|
type=int,
|
|
default=None,
|
|
help="最大处理文档数（用于测试）"
|
|
)
|
|
parser.add_argument(
|
|
"--no-verification",
|
|
action="store_true",
|
|
help="不使用二次对话验证（加快速度）"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# 加载配置
|
|
config = load_config()
|
|
|
|
# 确保输出目录存在
|
|
output_path = Path(args.output)
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
# 解析文档
|
|
logger.info(f"解析目录: {args.data_dir}")
|
|
parser_obj = DocumentParser()
|
|
docs = parser_obj.parse_directory(args.data_dir)
|
|
|
|
if args.max_docs:
|
|
docs = docs[:args.max_docs]
|
|
|
|
logger.info(f"找到 {len(docs)} 个文档")
|
|
|
|
# 切分为TextUnit
|
|
logger.info("切分文档为TextUnit...")
|
|
all_textunits = []
|
|
for doc in docs:
|
|
textunits = parser_obj.split_into_textunits(
|
|
doc,
|
|
max_length=config.MAX_TEXTUNIT_LENGTH
|
|
)
|
|
for tu in textunits:
|
|
tu["doc_id"] = doc.get("file_path", "")
|
|
tu["id"] = f"{doc.get('file_path', '')}_{tu.get('paragraph_index', 0)}"
|
|
all_textunits.extend(textunits)
|
|
|
|
logger.info(f"共生成 {len(all_textunits)} 个TextUnit")
|
|
|
|
# 构建知识图谱
|
|
logger.info("开始构建知识图谱...")
|
|
indexer = GraphIndexer(config=config)
|
|
kg_graph = indexer.index(
|
|
all_textunits,
|
|
use_verification=not args.no_verification
|
|
)
|
|
|
|
# 保存图谱
|
|
kg_obj = KnowledgeGraph()
|
|
kg_obj.graph = kg_graph
|
|
kg_obj.save(str(output_path), format="json")
|
|
|
|
# 保存统计信息
|
|
stats = {
|
|
"num_docs": len(docs),
|
|
"num_textunits": len(all_textunits),
|
|
"num_nodes": kg_graph.number_of_nodes(),
|
|
"num_edges": kg_graph.number_of_edges(),
|
|
"metadata": kg_graph.graph.get("metadata", {})
|
|
}
|
|
|
|
stats_path = output_path.parent / f"{output_path.stem}_stats.json"
|
|
with open(stats_path, "w", encoding="utf-8") as f:
|
|
json.dump(stats, f, ensure_ascii=False, indent=2)
|
|
|
|
logger.info(f"知识图谱构建完成！")
|
|
logger.info(f" - 节点数: {stats['num_nodes']}")
|
|
logger.info(f" - 边数: {stats['num_edges']}")
|
|
logger.info(f" - 结果已保存到: {output_path}")
|
|
logger.info(f" - 统计信息已保存到: {stats_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|
|
|
|
|
|
|