A simple Python library for building, pruning, evaluating and plotting colorful decision trees.
To install run pip install git+https://github.com/LisaAlaz/colortree@main
See example usage below.
Example usage:
import colortree as ct
import numpy as np
train = np.loadtxt('sample_data/train.txt')
dev = np.loadtxt('sample_data/dev.txt')
test = np.loadtxt('sample_data/test.txt')
# Build and plot tree from train set:
tree, depth = ct.make_tree(train)
# Build and plot optimally pruned tree from train set:
tree, depth = ct.make_pruned_tree(train, dev, n_classes=4)
# Inference on one sample:
pred = ct.predict_sample(tree, test[0])
# Inference on the test set:
test_preds = ct.predict_test_set(tree, test)
# Print evaluation report:
ct.eval_report(tree, test, n_classes=4)