xgboost_node_gradient_boosting_quickstart_train_predict.ts

1import * as xgboost from 'xgboost-node';
2
3async function main() {
4    // 1. Prepare your data (Features and Labels)
5    const trainData = [
6        [1, 2, 3],
7        [4, 5, 6],
8        [7, 8, 9]
9    ];
10    const trainLabels = [0, 0, 1];
11
12    // 2. Load data into DMatrix (the internal data structure for XGBoost)
13    const dtrain = new xgboost.DMatrix(trainData, trainLabels);
14
15    // 3. Set training parameters
16    const params = {
17        objective: 'binary:logistic',
18        eta: 0.1,
19        max_depth: 3
20    };
21
22    // 4. Train the model
23    const booster = await xgboost.train(params, dtrain, 10);
24
25    // 5. Predict
26    const testData = [[2, 3, 4]];
27    const dtest = new xgboost.DMatrix(testData);
28    const predictions = await booster.predict(dtest);
29
30    console.log('Predictions:', predictions);
31}
32
33main().catch(console.error);