Search in sources :

Example 1 with DecisionTreeRegressionModel

use of org.apache.spark.ml.regression.DecisionTreeRegressionModel 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;
}
Also used : DecisionTreeRegressionModel(org.apache.spark.ml.regression.DecisionTreeRegressionModel) Visitor(org.dmg.pmml.Visitor) Node(org.dmg.pmml.tree.Node) ClassifierNode(org.dmg.pmml.tree.ClassifierNode) BranchNode(org.dmg.pmml.tree.BranchNode) LeafNode(org.dmg.pmml.tree.LeafNode) DecisionTreeClassificationModel(org.apache.spark.ml.classification.DecisionTreeClassificationModel) TreeModelCompactor(org.jpmml.sparkml.visitors.TreeModelCompactor) TreeModel(org.dmg.pmml.tree.TreeModel) DecisionTreeModel(org.apache.spark.ml.tree.DecisionTreeModel) ImpurityCalculator(org.apache.spark.mllib.tree.impurity.ImpurityCalculator) CategoricalLabel(org.jpmml.converter.CategoricalLabel) LeafNode(org.dmg.pmml.tree.LeafNode) ClassifierNode(org.dmg.pmml.tree.ClassifierNode)

Aggregations

DecisionTreeClassificationModel (org.apache.spark.ml.classification.DecisionTreeClassificationModel)1 DecisionTreeRegressionModel (org.apache.spark.ml.regression.DecisionTreeRegressionModel)1 DecisionTreeModel (org.apache.spark.ml.tree.DecisionTreeModel)1 ImpurityCalculator (org.apache.spark.mllib.tree.impurity.ImpurityCalculator)1 Visitor (org.dmg.pmml.Visitor)1 BranchNode (org.dmg.pmml.tree.BranchNode)1 ClassifierNode (org.dmg.pmml.tree.ClassifierNode)1 LeafNode (org.dmg.pmml.tree.LeafNode)1 Node (org.dmg.pmml.tree.Node)1 TreeModel (org.dmg.pmml.tree.TreeModel)1 CategoricalLabel (org.jpmml.converter.CategoricalLabel)1 TreeModelCompactor (org.jpmml.sparkml.visitors.TreeModelCompactor)1