src/models/train_model.py

from catboost import CatBoostRegressor

def train_model(

        train,

        model_params,

        model_type,

        cat_features,

        valid=None,

):
    
    X_train, y_train = train
    
    if model_type == 'CATBOOST':

        model = CatBoostRegressor(**model_params,
                                  cat_features=cat_features)
        
        if valid:
            X_valid, y_valid = valid
            eval_set=[(X_valid,y_valid)]
        
        model.fit(X_train,y_train,
            eval_set=eval_set,
            verbose=200
        )

    return model