xgboost_node_binary_classification_train_and_predict.ts

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