62 lines
1.9 KiB
Python
62 lines
1.9 KiB
Python
#!/usr/bin/env python3
|
|
"""P&ID 펌프 추출기
|
|
|
|
P-10101, VP-10117, DP-10101 등 펌프/압축기 전용 추출.
|
|
|
|
사용법:
|
|
python pid_extract_pump.py --input full_text.txt --output pump.json
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
from pid_extract_template import call_llm
|
|
from pid_extract_prompts import _PUMP_PROMPT
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import time
|
|
|
|
logger = logging.getLogger("pid_extractor.pump")
|
|
|
|
def extract(input_text: str, max_tokens: int = 65536) -> list:
|
|
"""펌프/압축기 태그 추출."""
|
|
return call_llm(_PUMP_PROMPT, input_text, max_tokens=max_tokens)
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="P&ID 펌프 추출기")
|
|
parser.add_argument("--input", required=True, help="입력 텍스트 파일 경로")
|
|
parser.add_argument("--output", required=True, help="출력 JSON 파일 경로")
|
|
parser.add_argument("--max-tokens", type=int, default=65536, help="최대 토큰 수")
|
|
|
|
args = parser.parse_args()
|
|
|
|
with open(args.input, "r", encoding="utf-8") as f:
|
|
input_text = f.read()
|
|
|
|
logger.info(f"입력 파일 읽기 완료: {len(input_text)}자")
|
|
|
|
t0 = time.time()
|
|
tags = extract(input_text, max_tokens=args.max_tokens)
|
|
elapsed = time.time() - t0
|
|
|
|
logger.info(f"추출 완료: {len(tags)}개 태그, 소요 시간: {elapsed:.1f}초")
|
|
|
|
output_dir = os.path.dirname(args.output)
|
|
if output_dir:
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
result = {
|
|
"success": True,
|
|
"count": len(tags),
|
|
"tags": tags,
|
|
"processing_time_sec": round(elapsed, 1),
|
|
}
|
|
|
|
with open(args.output, "w", encoding="utf-8") as f:
|
|
json.dump(result, f, ensure_ascii=False, indent=2)
|
|
|
|
logger.info(f"결과 저장 완료: {args.output}")
|
|
print(json.dumps({"success": True, "count": len(tags), "time": round(elapsed, 1)}, ensure_ascii=False)) |