Skip to content

Commit a972668

Browse files
authored
docs: add snippet for creating boosted tree model (#1142)
* docs: create boosted tree model * merge main * update model * update test
1 parent 6fec681 commit a972668

File tree

1 file changed

+25
-1
lines changed

1 file changed

+25
-1
lines changed

samples/snippets/classification_boosted_tree_model_test.py

+25-1
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515

1616
def test_boosted_tree_model(random_model_id: str) -> None:
17-
# your_model_id = random_model_id
17+
your_model_id = random_model_id
1818
# [START bigquery_dataframes_bqml_boosted_tree_prepare]
1919
import bigframes.pandas as bpd
2020

@@ -39,4 +39,28 @@ def test_boosted_tree_model(random_model_id: str) -> None:
3939
)
4040
del input_data["functional_weight"]
4141
# [END bigquery_dataframes_bqml_boosted_tree_prepare]
42+
# [START bigquery_dataframes_bqml_boosted_tree_create]
43+
from bigframes.ml import ensemble
44+
45+
# input_data is defined in an earlier step.
46+
training_data = input_data[input_data["dataframe"] == "training"]
47+
X = training_data.drop(columns=["income_bracket", "dataframe"])
48+
y = training_data["income_bracket"]
49+
50+
# create and train the model
51+
census_model = ensemble.XGBClassifier(
52+
n_estimators=1,
53+
booster="gbtree",
54+
tree_method="hist",
55+
max_iterations=1, # For a more accurate model, try 50 iterations.
56+
subsample=0.85,
57+
)
58+
census_model.fit(X, y)
59+
60+
census_model.to_gbq(
61+
your_model_id, # For example: "your-project.census.census_model"
62+
replace=True,
63+
)
64+
# [END bigquery_dataframes_bqml_boosted_tree_create]
4265
assert input_data is not None
66+
assert census_model is not None

0 commit comments

Comments
 (0)