1  # Based on the official regression example:
 2  # https://catboost.ai/docs/concepts/python-usages-examples.html#regression
 3  
 4  import numpy as np
 5  from catboost import CatBoostRegressor
 6  
 7  import mlflow
 8  from mlflow.models import infer_signature
 9  
10  # Initialize data
11  train_data = np.array([[1, 4, 5, 6], [4, 5, 6, 7], [30, 40, 50, 60]])
12  train_labels = np.array([10, 20, 30])
13  eval_data = np.array([[2, 4, 6, 8], [1, 4, 50, 60]])
14  
15  # Initialize CatBoostRegressor
16  params = {
17      "iterations": 2,
18      "learning_rate": 1,
19      "depth": 2,
20      "allow_writing_files": False,
21  }
22  model = CatBoostRegressor(**params)
23  
24  # Fit model
25  model.fit(train_data, train_labels)
26  
27  # Log parameters and fitted model
28  with mlflow.start_run() as run:
29      signature = infer_signature(eval_data, model.predict(eval_data))
30      mlflow.log_params(params)
31      model_info = mlflow.catboost.log_model(model, name="model", signature=signature)
32  
33  # Load model
34  loaded_model = mlflow.catboost.load_model(model_info.model_uri)
35  
36  # Get predictions
37  preds = loaded_model.predict(eval_data)
38  print("predictions:", preds)