yggdrasil_ydf_random_forest_train_evaluate_predict.py

1import ydf
2import pandas as pd
3
4# 1. Load a dataset using Pandas
5ds_path = "https://raw.githubusercontent.com/google/yggdrasil-decision-forests/main/yggdrasil_decision_forests/test_data/dataset/adult_train.csv"
6train_ds = pd.read_csv(ds_path)
7
8# 2. Train a Random Forest model
9# The "label" argument specifies the column to predict
10model = ydf.RandomForestLearner(label="income").train(train_ds)
11
12# 3. Inspect the model (optional)
13print(model.summary())
14
15# 4. Evaluate the model
16evaluation = model.evaluate(train_ds)
17print(evaluation)
18
19# 5. Make predictions
20predictions = model.predict(train_ds.head())
21print(predictions)
22
23# 6. Save the model
24model.save("/tmp/my_model")