Examples with DecisionTreeClassificationModel org.apache.spark.ml.classification.DecisionTreeClassificationModel used on opensource projects

Example 1 with DecisionTreeClassificationModel

use of org.apache.spark.ml.classification.DecisionTreeClassificationModel in project jpmml-sparkml by jpmml.

the class TreeModelUtil method encodeDecisionTree.

private static <M extends Model<M> & DecisionTreeModel> TreeModel encodeDecisionTree(ModelConverter<?> converter, M model, PredicateManager predicateManager, ScoreDistributionManager scoreDistributionManager, Schema schema) {
    TreeModel treeModel;
    if (model instanceof DecisionTreeRegressionModel) {
        ScoreEncoder scoreEncoder = new ScoreEncoder() {

            @Override
            public Node encode(Node node, org.apache.spark.ml.tree.LeafNode leafNode) {
                node.setScore(leafNode.prediction());
                return node;
            }
        };
        treeModel = encodeTreeModel(MiningFunction.REGRESSION, scoreEncoder, model, predicateManager, schema);
    } else if (model instanceof DecisionTreeClassificationModel) {
        ScoreEncoder scoreEncoder = new ScoreEncoder() {

            private CategoricalLabel categoricalLabel = (CategoricalLabel) schema.getLabel();

            @Override
            public Node encode(Node node, org.apache.spark.ml.tree.LeafNode leafNode) {
                node = new ClassifierNode(null, node.getPredicate());
                int index = ValueUtil.asInt(leafNode.prediction());
                node.setScore(this.categoricalLabel.getValue(index));
                ImpurityCalculator impurityCalculator = leafNode.impurityStats();
                node.setRecordCount(ValueUtil.narrow(impurityCalculator.count()));
                scoreDistributionManager.addScoreDistributions(node, this.categoricalLabel.getValues(), impurityCalculator.stats());
                return node;
            }
        };
        treeModel = encodeTreeModel(MiningFunction.CLASSIFICATION, scoreEncoder, model, predicateManager, schema);
    } else {
        throw new IllegalArgumentException();
    }
    Boolean compact = (Boolean) converter.getOption(HasTreeOptions.OPTION_COMPACT, Boolean.TRUE);
    if (compact != null && compact) {
        Visitor visitor = new TreeModelCompactor();
        visitor.applyTo(treeModel);
    }
    return treeModel;
}