68 lines
1.9 KiB
Python
68 lines
1.9 KiB
Python
"""
|
|
모델 로드/저장/예측 담당
|
|
"""
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import joblib
|
|
import numpy as np
|
|
from loguru import logger
|
|
from sklearn.pipeline import Pipeline
|
|
|
|
import config
|
|
|
|
MODEL_FILENAME = "classifier.pkl"
|
|
|
|
|
|
def _get_latest_path() -> Optional[Path]:
|
|
candidates = sorted(config.CLASSIFIER_MODEL_BASE_PATH.rglob(MODEL_FILENAME))
|
|
return candidates[-1] if candidates else None
|
|
|
|
|
|
def _get_save_path() -> Path:
|
|
from datetime import datetime
|
|
date_dir = config.CLASSIFIER_MODEL_BASE_PATH / datetime.now().strftime("%Y%m%d")
|
|
date_dir.mkdir(exist_ok=True)
|
|
existing = sorted([d for d in date_dir.iterdir() if d.is_dir() and d.name.isdigit()])
|
|
next_idx = int(existing[-1].name) + 1 if existing else 1
|
|
save_dir = date_dir / f"{next_idx:04d}"
|
|
save_dir.mkdir(exist_ok=True)
|
|
return save_dir / MODEL_FILENAME
|
|
|
|
|
|
_model: Optional[Pipeline] = None
|
|
|
|
# 서버 시작 시 최신 모델 자동 로드
|
|
_latest = _get_latest_path()
|
|
if _latest:
|
|
try:
|
|
_model = joblib.load(_latest)
|
|
logger.info(f"classifier 모델 로드: {_latest} | classes={list(_model.classes_)}")
|
|
except Exception as e:
|
|
logger.warning(f"모델 로드 실패: {e}")
|
|
|
|
|
|
def get_model() -> Optional[Pipeline]:
|
|
return _model
|
|
|
|
|
|
def set_model(pipeline: Pipeline) -> None:
|
|
global _model
|
|
path = _get_save_path()
|
|
joblib.dump(pipeline, path)
|
|
_model = pipeline
|
|
logger.info(f"classifier 모델 저장: {path}")
|
|
|
|
|
|
def predict(
|
|
pipeline: Pipeline,
|
|
text: str,
|
|
threshold: float = 0.6,
|
|
) -> tuple[str, float, dict[str, float]]:
|
|
probs = pipeline.predict_proba([text])[0]
|
|
classes = pipeline.classes_
|
|
max_idx = int(np.argmax(probs))
|
|
confidence = float(probs[max_idx])
|
|
label = "OTHER" if confidence < threshold else classes[max_idx]
|
|
all_probs = {c: round(float(p), 4) for c, p in zip(classes, probs)}
|
|
return label, round(confidence, 4), all_probs |