97 lines
2.4 KiB
Python
97 lines
2.4 KiB
Python
#!/usr/bin/env python
|
||
"""
|
||
知识抽取脚本
|
||
示例用法：python scripts/extract.py --input data/sample.txt --output output/entities.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.extraction.ner import NERExtractor
|
||
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(
|
||
"--input",
|
||
type=str,
|
||
required=True,
|
||
help="输入文件路径（.docx或.txt）"
|
||
)
|
||
parser.add_argument(
|
||
"--output",
|
||
type=str,
|
||
default="output/entities.json",
|
||
help="输出文件路径"
|
||
)
|
||
parser.add_argument(
|
||
"--use-verification",
|
||
action="store_true",
|
||
default=True,
|
||
help="使用二次对话验证"
|
||
)
|
||
|
||
args = parser.parse_args()
|
||
|
||
# 加载配置
|
||
config = load_config()
|
||
|
||
# 确保输出目录存在
|
||
output_path = Path(args.output)
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
# 解析文档
|
||
input_path = Path(args.input)
|
||
if input_path.suffix == ".docx":
|
||
parser = DocumentParser()
|
||
doc_data = parser.parse_docx(str(input_path))
|
||
text = doc_data["text"]
|
||
else:
|
||
# 假设是纯文本文件
|
||
with open(input_path, "r", encoding="utf-8") as f:
|
||
text = f.read()
|
||
|
||
# 执行NER
|
||
logger.info("开始实体识别...")
|
||
ner_extractor = NERExtractor(config=config)
|
||
entities = ner_extractor.extract(text, use_verification=args.use_verification)
|
||
|
||
# 保存结果
|
||
result = {
|
||
"input_file": str(input_path),
|
||
"entities": entities,
|
||
"statistics": {
|
||
entity_type: len(entity_list)
|
||
for entity_type, entity_list in entities.items()
|
||
}
|
||
}
|
||
|
||
with open(output_path, "w", encoding="utf-8") as f:
|
||
json.dump(result, f, ensure_ascii=False, indent=2)
|
||
|
||
logger.info(f"实体识别完成，结果已保存到: {output_path}")
|
||
logger.info(f"统计信息: {result['statistics']}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
|
||
|
||
|
||
|
||
|