138 lines
4.3 KiB
Python
138 lines
4.3 KiB
Python
#!/usr/bin/env python
|
|
"""
|
|
完整流程脚本：从文档到知识图谱构建和查询
|
|
示例用法：python scripts/complete_pipeline.py --data-dir ../data/1法律 --output-dir ./output
|
|
"""
|
|
|
|
import argparse
|
|
import logging
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# 修复：移除conda环境的路径，确保使用虚拟环境的包
|
|
sys.path = [p for p in sys.path if 'anaconda3' not in p.lower() and 'conda' not in p.lower()]
|
|
# 添加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-dir",
|
|
type=str,
|
|
default="output",
|
|
help="输出目录"
|
|
)
|
|
parser.add_argument(
|
|
"--max-docs",
|
|
type=int,
|
|
default=None,
|
|
help="最大处理文档数（用于测试）"
|
|
)
|
|
parser.add_argument(
|
|
"--skip-index",
|
|
action="store_true",
|
|
help="跳过索引构建（使用已有索引）"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# 加载配置
|
|
config = load_config()
|
|
|
|
# 创建输出目录
|
|
output_dir = Path(args.output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
kg_file = output_dir / "knowledge_graph.json"
|
|
|
|
# 步骤1: 构建知识图谱索引
|
|
if not args.skip_index:
|
|
logger.info("=" * 60)
|
|
logger.info("步骤1: 解析文档并构建知识图谱")
|
|
logger.info("=" * 60)
|
|
|
|
# 解析文档
|
|
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=True)
|
|
|
|
# 保存图谱
|
|
kg_obj = KnowledgeGraph()
|
|
kg_obj.graph = kg_graph
|
|
kg_obj.save(str(kg_file), format="json")
|
|
|
|
logger.info(f"知识图谱已保存到: {kg_file}")
|
|
else:
|
|
logger.info(f"跳过索引构建，从文件加载: {kg_file}")
|
|
kg_obj = KnowledgeGraph()
|
|
kg_obj.load(str(kg_file), format="json")
|
|
|
|
# 步骤2: 展示图谱统计信息
|
|
logger.info("=" * 60)
|
|
logger.info("步骤2: 知识图谱统计信息")
|
|
logger.info("=" * 60)
|
|
|
|
if kg_obj.graph:
|
|
metadata = kg_obj.graph.graph.get("metadata", {})
|
|
logger.info(f"节点数: {kg_obj.graph.number_of_nodes()}")
|
|
logger.info(f"边数: {kg_obj.graph.number_of_edges()}")
|
|
logger.info(f"三元组数: {metadata.get('num_triplets', 0)}")
|
|
logger.info(f"社区数: {metadata.get('num_communities', 0)}")
|
|
|
|
# 展示一些社区摘要
|
|
summaries = metadata.get("summaries", {})
|
|
if summaries:
|
|
logger.info("\n前5个社区摘要：")
|
|
for i, (comm_id, comm_data) in enumerate(list(summaries.items())[:5], 1):
|
|
logger.info(f"{i}. {comm_id}: {comm_data.get('summary', '无摘要')[:100]}...")
|
|
|
|
logger.info("\n知识图谱构建流程完成！")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|