File size: 3,207 Bytes
101ed53 b72c967 101ed53 b72c967 101ed53 b72c967 101ed53 b72c967 101ed53 b72c967 101ed53 b72c967 101ed53 b72c967 101ed53 b72c967 101ed53 b72c967 101ed53 b72c967 101ed53 b72c967 101ed53 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | # modeling_catboost.py
import os
import json
import pickle
import numpy as np
import pandas as pd
from sklearn.preprocessing import LabelEncoder
class EssayCatBoostModel:
"""
Wrapper для сохранения/загрузки CatBoost модели с feature engineering.
Совместим с HuggingFace Hub (save_pretrained / from_pretrained).
"""
def __init__(self):
self.cat_models = []
self.task_classes = []
self.prompt_classes = []
self.feature_names = None
def save_pretrained(self, save_directory: str):
"""Сохраняет модель в папку для загрузки на HuggingFace."""
os.makedirs(save_directory, exist_ok=True)
# Сохраняем все CatBoost модели (по одной на фолд)
for i, model in enumerate(self.cat_models):
model_path = os.path.join(save_directory, f"catboost_fold_{i}.cbm")
model.save_model(model_path)
# Сохраняем метаданные
metadata = {
"n_folds": len(self.cat_models),
"task_classes": self.task_classes,
"prompt_classes": self.prompt_classes,
"feature_names": self.feature_names,
}
with open(os.path.join(save_directory, "metadata.json"), "w", encoding="utf-8") as f:
json.dump(metadata, f, ensure_ascii=False, indent=2)
print(f"CatBoost модель сохранена в {save_directory}")
@classmethod
def from_pretrained(cls, model_directory: str):
"""Загружает модель из папки (или с HuggingFace Hub)."""
import catboost as cb
instance = cls()
# Загружаем метаданные
metadata_path = os.path.join(model_directory, "metadata.json")
with open(metadata_path, "r", encoding="utf-8") as f:
metadata = json.load(f)
instance.task_classes = metadata.get("task_classes", [])
instance.prompt_classes = metadata.get("prompt_classes", [])
instance.feature_names = metadata.get("feature_names", None)
# Загружаем модели каждого фолда
n_folds = metadata["n_folds"]
for i in range(n_folds):
model_path = os.path.join(model_directory, f"catboost_fold_{i}.cbm")
model = cb.CatBoostRegressor()
model.load_model(model_path)
instance.cat_models.append(model)
print(f"CatBoost модель загружена из {model_directory} ({n_folds} фолдов)")
return instance
def predict(self, X: pd.DataFrame):
"""Предсказание усреднением по всем фолдам."""
if not self.cat_models:
raise ValueError("Модель не загружена")
# Убеждаемся, что порядок колонок совпадает с обучением
if self.feature_names is not None:
X = X[self.feature_names]
preds = np.zeros(len(X))
for model in self.cat_models:
preds += model.predict(X) / len(self.cat_models)
return preds |