train_model.py - master - maruga/mlops

train_model.py
335 строк12 KB
s546371
debug clearml pipe
28 мар 2026, 20:45
28 мар 2026, 20:459f709f8
from pathlib import Path
from typing import Any, Dict, List, Tuple, Union

import joblib
import numpy as np
import optuna
import pandas as pd
import s3fs
from catboost import CatBoostClassifier
from sklearn.metrics import average_precision_score
from sklearn.model_selection import train_test_split

from src.config import (
    CatBoostClsConfig,
    DataMap,
    S3Config,
    TrainConfig,
    config_s3,
)


def read_data(
    data_path: Path | str,
    storage_options: Dict[str, Union[str, Dict[str, str]]] | None = None,
) -> pd.DataFrame:
    """Читает данные .parquet по пути из конфигурации, где может быть
    указан доступ до S3 или локальный путь
    Returns:
        структурированные данных pandas
    """
    if str(data_path).startswith(("s3://", "s3a://")):
        if not storage_options:
            raise ValueError("Требуются ключи доступа к S3")

        if "client_kwargs" not in storage_options:
            raise KeyError("client_kwargs отсутствует в storage_options")
        client_kwargs = storage_options["client_kwargs"]
        if not isinstance(client_kwargs, dict):
            raise TypeError("client_kwargs должен быть словарём")
        endpoint_url: str = client_kwargs["endpoint_url"]

        if "endpoint_url" not in client_kwargs:
            raise KeyError("endpoint_url отсутствует в client_kwargs")
        simplified_storage_options = {
            "key": storage_options["key"],
            "secret": storage_options["secret"],
            "endpoint_url": endpoint_url,
        }
        storage_options = simplified_storage_options
    else:
        storage_options = None

    try:
        df = pd.read_parquet(data_path, storage_options=storage_options)
        return df
    except Exception as e:
        raise Exception(f"Ошибка чтения {data_path}: {e}") from e


def prepare_data(
    df: pd.DataFrame,
    features: List[str] = DataMap.features,
    target: str = DataMap.target,
    test_size: float = TrainConfig.TEST_SIZE,
    random_state: int = TrainConfig.RANDOM_STATE,
) -> Tuple[pd.DataFrame, pd.DataFrame, pd.Series, pd.Series] | Any:
    """
    Подготавливает данные
    Args:
        df: структурированные данных  pandas
        features: список признаков
        target: целевой признак
        test_size: размер тестовой выборки
        random_state: значение для воспроизводимости случайных значений,
    Returns:
        кортеж с X_train, X_test, y_train, y_test
    Raises:
        ValueError: если список признаков или целевой признак пусты
        KeyError: если указанная колонка не найдена в данных
        Exception: при других непредвиденных ошибках
    """
    if not features or not target:
        raise ValueError("Список признаков и целевой признак не должны быть пустыми")

    try:
        X = df[features].astype(float)
        y = df[target]
        splitted_data = train_test_split(
            X, y, test_size=test_size, random_state=random_state, stratify=y
        )
        return splitted_data
    except KeyError as e:
        raise KeyError(f"Колонка не найдена в данных: {e}") from e
    except Exception as e:
        raise Exception(f"Непредвиденная ошибка: {e}") from e


def objective(
    trial: optuna.trial.Trial,
    X_train: pd.DataFrame,
    y_train: pd.Series,
    X_val: pd.DataFrame,
    y_val: pd.Series,
) -> Any:
    """Целевая функция для Optuna
    Args:
        trial:
        X_train: обучающая выборка
        y_train: зависимая переменная
        X_val: набор для проверки
        y_val: переменная для проверки
    Returns:
        Средний балл точности - PR-AUC
    """
    params_range = {
        "depth": trial.suggest_int("depth", 3, 10),
        "learning_rate": trial.suggest_float("learning_rate", 0.01, 0.3, log=True),
        "iterations": trial.suggest_int("iterations", 100, 1000),
        "l2_leaf_reg": trial.suggest_float("l2_leaf_reg", 1e-5, 10, log=True),
        "border_count": trial.suggest_int("border_count", 32, 255),
        "random_strength": trial.suggest_float("random_strength", 1e-5, 10, log=True),
        "auto_class_weights": "SqrtBalanced",
        "verbose": False,
        "random_seed": TrainConfig.RANDOM_STATE,
    }

    trial_model = CatBoostClassifier(**params_range)

    trial_model.fit(
        X_train,
        y_train,
        eval_set=(X_val, y_val),
        early_stopping_rounds=50,
        verbose=False,
    )

    y_pred_proba = trial_model.predict_proba(X_val)[:, 1]
    pr_auc = average_precision_score(y_val, y_pred_proba)

    return pr_auc


def train_with_optuna(
    X_train: pd.DataFrame,
    y_train: pd.Series,
    test_size: float = TrainConfig.TEST_SIZE,
    random_state: int = TrainConfig.RANDOM_STATE,
    n_trials: int = 50,
) -> Dict[str, Union[float, int]]:
    """
    Подбор гиперпараметров с OptunaArgs
    Args:
        X_train: обучающая выборка
        y_train: зависимая переменная
        test_size: размер тестовой выборки
        random_state: величина для воспроизводимости случайных значений
        n_trials: число подходов, по умолчанию 50
    Returns:
        Словарь с лучшими гиперпараметрами
    """
    X_train_sub, X_val, y_train_sub, y_val = train_test_split(
        X_train,
        y_train,
        test_size=test_size,
        random_state=random_state,
        stratify=y_train,
    )
    study = optuna.create_study(
        direction="maximize",
        sampler=optuna.samplers.TPESampler(seed=random_state),
    )
    study.optimize(
        lambda trial: objective(trial, X_train_sub, y_train_sub, X_val, y_val),
        n_trials=n_trials,
        show_progress_bar=True,
    )
    print(f"Лучшее значение PR-AUC: {study.best_value:.4f}")
    print(f"Лучшие гиперпараметры: {study.best_params}")

    return study.best_params


def train_boost_model(
    X_train: pd.DataFrame,
    y_train: pd.Series,
    params_model: Dict[str, Union[float, int]],
    random_state: int = TrainConfig.RANDOM_STATE,
    get_optuna: bool = False,
) -> CatBoostClassifier:
    """
    Обучает модель
    Args:
        X_train: обучающая выборка
        y_train: зависимая переменная
        params_model: гиперпараметры для обучения модели
        random_state: величина для воспроизводимости случайных значений
        get_optuna: использовать Optuna для подбора параметров
    Returns:
        CatBoostClassifier
    """
    if get_optuna:
        print("\nПодбор гиперпараметров с Optuna...")
        params_model = train_with_optuna(
            X_train, y_train, n_trials=30, random_state=random_state
        )

    model = CatBoostClassifier(
        **params_model,
        auto_class_weights="SqrtBalanced",
        random_seed=random_state,
        verbose=False,
    )
    model.fit(X_train, y_train)
    return model


def validate_model(
    model: CatBoostClassifier, X_test: pd.DataFrame, y_test: pd.Series
) -> Dict[str, float]:
    """
    Определяет метрики модели по тестовой выборке
    Args:
        model: модель для валидации
        X_test: тестовая выборка
        y_test: метки категорий
    Returns:
        Словарь с метриками и их значением
    """
    y_pred = model.predict(X_test)
    y_pred_proba = model.predict_proba(X_test)[:, 1]

    from sklearn.metrics import (
        accuracy_score,
        average_precision_score,
        confusion_matrix,
        f1_score,
        mean_absolute_error,
        mean_squared_error,
        precision_score,
        r2_score,
        recall_score,
    )

    metrics = {
        "mse": np.round(mean_squared_error(y_test, y_pred), 5).item(),
        "rmse": float(np.sqrt(mean_squared_error(y_test, y_pred))),
        "mae": float(mean_absolute_error(y_test, y_pred)),
        "r2": float(r2_score(y_test, y_pred)),
        "accuracy": float(accuracy_score(y_test, y_pred)),
        "precision": float(precision_score(y_test, y_pred, zero_division=0)),
        "recall": float(recall_score(y_test, y_pred, zero_division=0)),
        "f1": float(f1_score(y_test, y_pred, zero_division=0)),
        "pr_auc": float(average_precision_score(y_test, y_pred_proba)),
    }

    cm = confusion_matrix(y_test, y_pred)
    print(f"TN: {cm[0, 0]}, FP: {cm[0, 1]}, FN: {cm[1, 0]}, TP: {cm[1, 1]}")

    return metrics


def save_model(model: CatBoostClassifier, model_path: str) -> None:
    """
    Сохраняет модель в S3
    """
    if config_s3.storage_options is None:
        raise ValueError("storage_options не инициализирован")

    if "client_kwargs" not in config_s3.storage_options:
        raise KeyError("client_kwargs отсутствует в storage_options")

    client_kwargs = config_s3.storage_options["client_kwargs"]
    if not isinstance(client_kwargs, dict):
        raise TypeError("client_kwargs должен быть словарём")

    fs = s3fs.S3FileSystem(client_kwargs=config_s3.storage_options["client_kwargs"])
    with fs.open(model_path, "wb") as f:
        joblib.dump(model, f, compress=True)
    print(f"model saved to S3: {model_path}")


if __name__ == "__main__":

    def main() -> None:
        """
        Обучения модели
        """
        from dotenv import load_dotenv

        load_dotenv()

        print("Загрузка данных...")
        df = read_data(
            f"{config_s3.S3_DATA_PATH}/{TrainConfig.DATA_PATH}",
            config_s3.storage_options,
        )

        print("Подготовка данных...")
        X_train, X_test, y_train, y_test = prepare_data(df)

        params_model = {
            "depth": CatBoostClsConfig.DEPTH,
            "learning_rate": CatBoostClsConfig.LEARNING_RATE,
            "iterations": CatBoostClsConfig.ITERATIONS,
            "l2_leaf_reg": CatBoostClsConfig.L2_LEAF_REG,
            "border_count": CatBoostClsConfig.BORDER_COUNT,
            "random_strength": CatBoostClsConfig.RANDOM_STRENGTH,
        }
        """
        Confusion Matrix:
        TN: 517870, FP: 9124, FN: 4628, TP: 2500
        mse: 0.0257
        rmse: 0.1605
        mae: 0.0257
        r2: -0.9554
        accuracy: 0.9743
        precision: 0.2151
        recall: 0.3507
        f1: 0.2666
        pr_auc: 0.2094
        """

        print("\nОбучение модели...")
        model = train_boost_model(X_train, y_train, params_model)
        pos_ratio = (y_train == 1).mean() * 100
        print(f"Доля положительных классов в train: {pos_ratio:.2f}%")

        print("\nВалидация модели...")
        metrics = validate_model(model, X_test, y_test)
        for metric, value in metrics.items():
            print(f"{metric}: {value:.4f}")

        save_model(model, f"{S3Config.S3_MODEL_PATH}/{TrainConfig.MODEL_PATH}")

    main()