""" NLP 통합 엔진 구두점 복원 및 MLM 예측을 통합 관리하는 엔진입니다. """ import os import json import torch import logging from pathlib import Path from typing import Dict, Any, Optional, List from .nlp.punctuation_restorer import PunctuationRestorer from .nlp.mlm_predictor import MLMPredictor from .nlp.utils import remove_punctuation, replace_mask_with_symbol logger = logging.getLogger(__name__) def load_nlp_config(config_path: Optional[str] = None) -> Dict[str, Any]: """ NLP 설정 파일을 로드합니다. Args: config_path: 설정 파일 경로 (None이면 기본 경로 사용) Returns: 설정 딕셔너리 """ if config_path is None: config_path = Path(__file__).parent / "config" / "nlp_config.json" else: config_path = Path(config_path) if not config_path.exists(): raise FileNotFoundError(f"NLP 설정 파일을 찾을 수 없습니다: {config_path}") with open(config_path, 'r', encoding='utf-8') as f: return json.load(f) class NLPEngine: """NLP 처리 통합 엔진 클래스""" def __init__(self, config_path: Optional[str] = None): """ NLP 엔진을 초기화합니다. Args: config_path: 설정 파일 경로 (None이면 기본 경로 사용) """ self.config = load_nlp_config(config_path) # 디바이스 설정 dev_cfg = self.config.get('device', 'auto') if dev_cfg == 'auto': self.device = "cuda" if torch.cuda.is_available() else "cpu" else: self.device = dev_cfg logger.info(f"[NLP] Device: {self.device}") # 모델 캐시 경로 (환경 변수 또는 기본값) self.base_model_dir = os.getenv( 'AI_MODEL_DIR', str(Path(__file__).parent.parent / "models") ) # 서브 모듈 초기화 (지연 로딩) self.punc_restorer = None self.mlm_predictor = None def _load_models(self): """필요할 때 모델을 메모리에 로드""" if self.punc_restorer is None: logger.info("[NLP] 구두점 복원 모델 로드 중...") self.punc_restorer = PunctuationRestorer( self.config, self.base_model_dir, self.device ) self.punc_restorer.download_model() self.punc_restorer.load_model() if self.mlm_predictor is None: logger.info("[NLP] MLM 모델 로드 중...") self.mlm_predictor = MLMPredictor(self.config, self.device) self.mlm_predictor.load_model() def process_text( self, raw_text: str, ocr_results: Optional[List[Dict]] = None, add_space: bool = True, reduce_punc: bool = True ) -> Dict[str, Any]: """ 텍스트 처리 파이프라인: 1. 구두점 제거 (전처리) 2. 구두점 복원 3. [MASK] 예측 Args: raw_text: 원본 텍스트 (구두점 포함 가능) add_space: 구두점 뒤 공백 추가 여부 reduce_punc: 구두점 단순화 여부 Returns: 처리 결과 딕셔너리 """ self._load_models() try: # 1. 전처리 (구두점 제거, [MASK] 보존) clean_text = remove_punctuation(raw_text) logger.info(f"[NLP] 구두점 제거 완료: {len(clean_text)} 글자") # 2. 구두점 복원 punctuated_text = self.punc_restorer.restore_punctuation( clean_text, add_space=add_space, reduce=reduce_punc ) logger.info(f"[NLP] 구두점 복원 완료: {len(punctuated_text)} 글자") # 3. MLM 예측 mask_predictions = self.mlm_predictor.predict_masks(punctuated_text) logger.info(f"[NLP] MLM 예측 완료: {len(mask_predictions)}개 마스크") # 4. 출력용 텍스트 생성 ([MASK] -> □) mask_replacement = self.config['tokens']['mask_replacement'] final_text = replace_mask_with_symbol( punctuated_text, mask_replacement ) # Extract mask info from OCR results or original text mask_info_list = [] if ocr_results: # Use OCR results to get order and type for item in ocr_results: if 'MASK' in item.get('type', ''): mask_info_list.append({ 'order': item.get('order', 0), 'type': item.get('type', 'MASK2'), 'text': item.get('text', '') }) else: # Fallback: extract from text i = 0 while i < len(raw_text): if raw_text[i] == '[' and 'MASK' in raw_text[i:i+10]: end = raw_text.find(']', i) if end != -1: mask_text = raw_text[i:end+1] mask_type = 'MASK1' if 'MASK1' in mask_text else 'MASK2' mask_info_list.append({ 'order': len(mask_info_list), # Sequential order 'type': mask_type, 'text': mask_text }) i = end + 1 continue i += 1 # Format results according to specification formatted_results = [] for idx, pred_list in enumerate(mask_predictions): if idx < len(mask_info_list): mask_info = mask_info_list[idx] formatted_results.append({ "order": mask_info['order'], "type": mask_info['type'], "top_10": pred_list[:10] # Top-10 predictions }) else: # Fallback if mask_info_list is shorter formatted_results.append({ "order": idx, "type": "MASK2", "top_10": pred_list[:10] }) # Calculate statistics top1_probs = [preds[0]['probability'] for preds in mask_predictions if preds] statistics = { "top1_probability_avg": float(sum(top1_probs) / len(top1_probs)) if top1_probs else 0.0, "top1_probability_min": float(min(top1_probs)) if top1_probs else 0.0, "top1_probability_max": float(max(top1_probs)) if top1_probs else 0.0, "total_masks": len(mask_predictions) } return { "punctuated_text_with_masks": final_text, "results": formatted_results, "statistics": statistics } except Exception as e: logger.error(f"[NLP] 처리 중 오류: {e}", exc_info=True) return { "success": False, "error": str(e) } def restore_punctuation_only( self, text: str, add_space: bool = True, reduce_punc: bool = True ) -> Dict[str, Any]: """ 구두점 복원만 수행합니다 (MLM 예측 제외). Args: text: 입력 텍스트 add_space: 구두점 뒤 공백 추가 여부 reduce_punc: 구두점 단순화 여부 Returns: 구두점 복원 결과 """ self._load_models() try: clean_text = remove_punctuation(text) punctuated_text = self.punc_restorer.restore_punctuation( clean_text, add_space=add_space, reduce=reduce_punc ) return { "success": True, "original_text": text, "clean_text": clean_text, "punctuated_text": punctuated_text } except Exception as e: logger.error(f"[NLP] 구두점 복원 중 오류: {e}", exc_info=True) return { "success": False, "error": str(e) } def predict_masks_only( self, text: str ) -> Dict[str, Any]: """ MLM 예측만 수행합니다 (구두점 복원 제외). Args: text: 마스크가 포함된 텍스트 Returns: MLM 예측 결과 """ self._load_models() try: mask_predictions = self.mlm_predictor.predict_masks(text) return { "success": True, "predictions": mask_predictions, "mask_count": len(mask_predictions) } except Exception as e: logger.error(f"[NLP] MLM 예측 중 오류: {e}", exc_info=True) return { "success": False, "error": str(e) } # ================================================================================ # Global Accessor # ================================================================================ _nlp_engine = None def get_nlp_engine(config_path: Optional[str] = None) -> NLPEngine: """ 전역 NLP 엔진 인스턴스를 반환합니다 (싱글톤 패턴). Args: config_path: 설정 파일 경로 (None이면 기본 경로 사용) Returns: NLPEngine 인스턴스 """ global _nlp_engine if _nlp_engine is None: _nlp_engine = NLPEngine(config_path) return _nlp_engine def process_text_with_nlp( text: str, ocr_results: Optional[List[Dict]] = None, config_path: Optional[str] = None, add_space: bool = True, reduce_punc: bool = True ) -> Dict[str, Any]: """ 편의 함수: 텍스트를 NLP 파이프라인으로 처리합니다. Args: text: 입력 텍스트 ocr_results: OCR 결과 리스트 (order, type 정보 포함) config_path: 설정 파일 경로 add_space: 구두점 뒤 공백 추가 여부 reduce_punc: 구두점 단순화 여부 Returns: 처리 결과 딕셔너리 """ engine = get_nlp_engine(config_path) return engine.process_text(text, ocr_results=ocr_results, add_space=add_space, reduce_punc=reduce_punc)