Train the GBT

For the gradient boosted trees model (GBT), the formatting of the training and evaluating dataset change a little, instead of provide numpy arrays, you should pass dictionnaries:

def run_gbt():
    (train_inputs, train_targets, train_timestamps), (test_inputs, test_targets, test_timestamps) = split_data()

    # creating datasets for training and testing
    train_dataset = tf.data.Dataset.from_tensor_slices(
        (
            {name: train_inputs[:, i] for i, name in enumerate(INPUTS_VAR)},
            train_targets[:, 0],
        ),
    ).batch(250)

    test_dataset = tf.data.Dataset.from_tensor_slices(
        (
            {name: test_inputs[:, i] for i, name in enumerate(INPUTS_VAR)},
            test_targets[:, 0],
        ),
    ).batch(250)

Moreover, you can automatically adjust the hyperparameters by passing a tuner to the model generator:

# creating a gradient boosted model
tuner = tfdf.tuner.RandomSearch(num_trials=250, use_predefined_hps=True)
model_set = nn.ModelSet(
    nb_models=1,
    model_generator=nn.gbt_model,
    tuner=tuner,
)

The following steps are identical to these with an ANN:

model_set.fit(train_dataset)
model_set.evaluate(test_dataset)
prediction = np.median(model_set(test_dataset)[:, 0], axis=0)

utils.write_output(prediction, column=TARGETS_VAR)