dagshub_project/code/train_model.py

import pandas as pd
import numpy as np
from catboost import CatBoostClassifier, Pool
from dagshub import dagshub_logger
from catboost import CatBoostClassifier
from sklearn.metrics import fbeta_score, precision_recall_curve
from catboost.utils import get_roc_curve, select_threshold
import pickle

def eval_model(model):
    with dagshub_logger() as logger:
        logger.log_hyperparams(model_class=type(model).__name__)
        logger.log_hyperparams({'model': model.get_params()})
        model.fit(X_train, y_train)
        train_preds = model.predict(X_train)
        test_preds = model.predict(X_val)
        logger.log_metrics(
            {'train__f2': fbeta_score(y_train, train_preds, beta=2), 'val_f2': fbeta_score(y_val, test_preds, beta=2)})
X_train = np.load('processed_data\\X_train_processed.npy', allow_pickle=True)
X_val = np.load('processed_data\\X_val_processed.npy', allow_pickle=True)
y_train = np.load('processed_data\\y_train_processed.npy', allow_pickle=True)
y_val = np.load('processed_data\\y_val_processed.npy', allow_pickle=True)
model =  CatBoostClassifier(iterations=200, cat_features=[0,1,2,3,4,5], random_state=0, scale_pos_weight=13.791)
eval_model(model)
with open('trained_models\\catboost.pk1', 'wb') as f:
    pickle.dump(model,f)