58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
#!/usr/bin/env python
|
|
"""
|
|
快速测试：只处理1个文档的前2个TextUnit
|
|
"""
|
|
|
|
import sys
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
sys.path = [p for p in sys.path if 'anaconda3' not in p.lower() and 'conda' not in p.lower()]
|
|
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.utils.config import load_config
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def main():
|
|
logger.info("=" * 60)
|
|
logger.info("快速测试：处理1个文档的前2个TextUnit")
|
|
logger.info("=" * 60)
|
|
|
|
config = load_config()
|
|
|
|
# 解析1个文档
|
|
parser = DocumentParser()
|
|
doc_path = "../data/1法律/10-中华人民共和国乡村振兴促进法.docx"
|
|
doc = parser.parse_docx(doc_path)
|
|
|
|
# 切分为TextUnit，只取前2个
|
|
textunits = parser.split_into_textunits(doc, max_length=config.MAX_TEXTUNIT_LENGTH)[:2]
|
|
for tu in textunits:
|
|
tu["doc_id"] = doc.get("file_path", "")
|
|
tu["id"] = f"{doc.get('file_path', '')}_{tu.get('paragraph_index', 0)}"
|
|
|
|
logger.info(f"准备处理 {len(textunits)} 个TextUnit")
|
|
|
|
# 构建知识图谱（只处理这2个）
|
|
logger.info("开始构建知识图谱...")
|
|
indexer = GraphIndexer(config=config)
|
|
kg_graph = indexer.index(textunits, use_verification=True)
|
|
|
|
logger.info("=" * 60)
|
|
logger.info("快速测试完成！")
|
|
logger.info(f"节点数: {kg_graph.number_of_nodes()}")
|
|
logger.info(f"边数: {kg_graph.number_of_edges()}")
|
|
logger.info("=" * 60)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|
|
|