feat: auto-download CatBoost models and on_topic from HF Space

src/predict.py

@@ -12,6 +12,7 @@ import numpy as np
 import pandas as pd
 import joblib
 from catboost import CatBoostRegressor
+from huggingface_hub import hf_hub_download  # <— автодозагрузка файлов из HF
 
 # --- импорты проекта ---
 HERE = Path(__file__).parent
@@ -26,10 +27,13 @@ except ModuleNotFoundError:
     # если файл лежит в src/
     from src.feature_engineering import FeatureExtractor  # type: ignore
 
-# --- пути ---
+# --- пути/константы ---
 MODELS_DIR = ROOT / "models"                   # catboost_Q1.cbm ... catboost_Q4.cbm
 ON_TOPIC_PATH = MODELS_DIR / "on_topic.pkl"    # опционально
 
+# репозиторий Space, откуда подтягиваем артефакты, если их нет локально
+SPACE_REPO = os.environ.get("SPACE_REPO", "lidiiakarmanova/exam-evaluator")
+
 # --- служебные колонки (не подавать в модель) ---
 NON_NUMERIC_KEEP = {"question_number", "question_text", "answer_text"}
 TARGET_COLS = {"score", "Оценка экзаменатора"}
@@ -75,11 +79,55 @@ def _clip_by_q(qnum: int, preds: np.ndarray) -> np.ndarray:
     return np.clip(preds, lo, hi)
 
 
+def _ensure_model_file(qnum: int) -> Path:
+    """
+    Гарантирует наличие файла модели Q{qnum} локально.
+    Если файла нет — скачивает из Space (путь в репо: models/catboost_Q{q}.cbm).
+    """
+    MODELS_DIR.mkdir(parents=True, exist_ok=True)
+    local_path = MODELS_DIR / f"catboost_Q{qnum}.cbm"
+    if local_path.exists():
+        return local_path
+
+    remote_filename = f"models/catboost_Q{qnum}.cbm"
+    print(f"[i] Модель Q{qnum} не найдена локально, скачиваем из {SPACE_REPO}:{remote_filename}")
+    cache_path = hf_hub_download(
+        repo_id=SPACE_REPO,
+        repo_type="space",
+        filename=remote_filename,
+    )
+    # скопируем из кэша в models/ — Space может чистить кэш между рестартами
+    Path(local_path).write_bytes(Path(cache_path).read_bytes())
+    return local_path
+
+
+def _ensure_on_topic_file() -> Path | None:
+    """
+    Если используем on_topic.pkl — аналогично подтянем из Space (models/on_topic.pkl),
+    иначе вернём None.
+    """
+    if ON_TOPIC_PATH.exists():
+        return ON_TOPIC_PATH
+
+    remote_filename = "models/on_topic.pkl"
+    try:
+        print(f"[i] on_topic.pkl не найден локально, пробуем скачать из {SPACE_REPO}:{remote_filename}")
+        cache_path = hf_hub_download(
+            repo_id=SPACE_REPO,
+            repo_type="space",
+            filename=remote_filename,
+        )
+        ON_TOPIC_PATH.parent.mkdir(parents=True, exist_ok=True)
+        Path(ON_TOPIC_PATH).write_bytes(Path(cache_path).read_bytes())
+        return ON_TOPIC_PATH
+    except Exception as e:
+        print(f"[!] Не удалось скачать on_topic.pkl: {e}")
+        return None
+
+
 def _load_model(qnum: int) -> CatBoostRegressor:
-    """Загрузка CatBoost-модели для указанного вопроса."""
-    model_path = MODELS_DIR / f"catboost_Q{qnum}.cbm"
-    if not model_path.exists():
-        raise FileNotFoundError(f"Не найден файл модели: {model_path}")
+    """Загрузка CatBoost-модели для указанного вопроса (с автодозагрузкой из HF)."""
+    model_path = _ensure_model_file(qnum)
     model = CatBoostRegressor()
     model.load_model(str(model_path))
     return model
@@ -98,10 +146,15 @@ def _align_to_model_features(model: CatBoostRegressor, X: pd.DataFrame) -> pd.Da
 
 def _maybe_add_on_topic(df_feats: pd.DataFrame) -> pd.DataFrame:
     """
-    Если есть on_topic.pkl (pack = {'model': clf, 'features': [...]})
+    Если есть on_topic.pkl (pack = {'model': clf, 'features': [...]}),
     — добавляем вероятность 'on_topic_prob'. Иначе 0.0.
     """
     out = df_feats.copy()
+
+    # попытаемся подтянуть on_topic.pkl из Space при необходимости
+    if not ON_TOPIC_PATH.exists():
+        _ensure_on_topic_file()
+
     if not ON_TOPIC_PATH.exists():
         out["on_topic_prob"] = 0.0
         return out
@@ -180,7 +233,7 @@ def pipeline_infer(input_csv: Path, output_csv: Path) -> None:
         pq = np.asarray(pq, dtype=float).reshape(-1)
         preds[mask.values] = _clip_by_q(q, pq)
 
-    # --- новое надёжное округление (без .loc по индексам) ---
+    # --- надёжное округление ---
     qnums = feats["question_number"].astype(int).to_numpy()
     rounded = np.rint(preds).astype(np.float32)
     mask13 = (qnums == 1) | (qnums == 3)