Initial commit

utils/classifier_model_store.py

@@ -0,0 +1,68 @@
+"""
+모델 로드/저장/예측 담당
+"""
+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