catboost_classifier_titanic_dataset_train_and_predict.py

1from catboost import CatBoostClassifier, Pool
2from catboost.datasets import titanic
3import numpy as np
4
5# Load training and testing data
6train_df, test_df = titanic()
7
8# Fill missing values
9train_df.fillna(-999, inplace=True)
10test_df.fillna(-999, inplace=True)
11
12# Separate features and target
13X = train_df.drop('Survived', axis=1)
14y = train_df.Survived
15
16# Identify categorical features indices
17cat_features = np.where(X.dtypes != float)[0]
18
19# Initialize CatBoostClassifier
20model = CatBoostClassifier(
21    iterations=10,
22    learning_rate=1,
23    depth=2
24)
25
26# Fit model
27model.fit(X, y, cat_features)
28
29# Get predicted classes
30preds_class = model.predict(test_df)
31
32# Get predicted probabilities
33preds_proba = model.predict_proba(test_df)
34
35# Get raw score
36preds_raw = model.predict(test_df, prediction_type='RawFormulaVal')
37
38print(f"Predictions: {preds_class[:5]}")