Search in sources :

Example 1 with GBTRegressionModel

use of org.apache.spark.ml.regression.GBTRegressionModel in project jpmml-sparkml by jpmml.

the class GBTRegressionModelConverter method encodeModel.

@Override
public MiningModel encodeModel(Schema schema) {
    GBTRegressionModel model = getTransformer();
    List<TreeModel> treeModels = TreeModelUtil.encodeDecisionTreeEnsemble(model, schema);
    MiningModel miningModel = new MiningModel(MiningFunction.REGRESSION, ModelUtil.createMiningSchema(schema.getLabel())).setSegmentation(MiningModelUtil.createSegmentation(Segmentation.MultipleModelMethod.WEIGHTED_SUM, treeModels, Doubles.asList(model.treeWeights())));
    return miningModel;
}
Also used : TreeModel(org.dmg.pmml.tree.TreeModel) MiningModel(org.dmg.pmml.mining.MiningModel) GBTRegressionModel(org.apache.spark.ml.regression.GBTRegressionModel)

Aggregations

GBTRegressionModel (org.apache.spark.ml.regression.GBTRegressionModel)1 MiningModel (org.dmg.pmml.mining.MiningModel)1 TreeModel (org.dmg.pmml.tree.TreeModel)1