@__username/decision-tree/examples/using-uci-datasets.js

NodeJS implementation of decision tree using ID3 algorithm

var DecisionTree = require('..');
var fs = require('fs');

var features = ["age", "workclass", "fnlwgt", "education", "education-num", "marital-status", "occupation", "relationship", "race", "sex", "capital-gain", "capital-loss", "hours-per-week", "native-country"];
var class_name = "result";

var columnMapping = features.concat(class_name);

function readData(file)
{
  var rows = fs.readFileSync(file)
    .toString()
    .split('\n')
    .slice(0, 100)
    .filter(function (l) {
      return l.trim().length > 0;
    })
    .map(function (row) {
      var d = {};
      var columns = row.split(', ');
      for (var i = 0; i < columns.length; ++i) {
        d[columnMapping[i]] = columns[i];
      }
      return d;
    });

  return rows;
}

var training_data = readData(__dirname + '/../data/uci-ml/adult/adult.data');
var test_data = readData(__dirname + '/../data/uci-ml/adult/adult.test');

console.log('Building tree...');
console.log('Training sample:\n', JSON.stringify(training_data.slice(0, 5), null, '  '));
console.log('Testing sample:\n', JSON.stringify(test_data.slice(0, 5), null, '  '));

var dt = new DecisionTree(class_name, features, {
  verbose: true
}).train(training_data);

console.log('Evaluating...');
console.log('Accuracy:');
var accuracy = dt.evaluate(test_data);
console.log(accuracy);

console.log('Exporting...');
var treeModel = dt.toJSON();
fs.writeFileSync('./model.json', treeModel);