diff --git a/pmml-converter/src/main/java/org/jpmml/converter/mining/MiningModelUtil.java b/pmml-converter/src/main/java/org/jpmml/converter/mining/MiningModelUtil.java index 67075b35..0d25635b 100644 --- a/pmml-converter/src/main/java/org/jpmml/converter/mining/MiningModelUtil.java +++ b/pmml-converter/src/main/java/org/jpmml/converter/mining/MiningModelUtil.java @@ -175,11 +175,8 @@ public MiningModel createModelChain(List models, Segmentation.M MiningFunction miningFunction = lastModel.requireMiningFunction(); - Segmentation segmentation = createSegmentation(Segmentation.MultipleModelMethod.MODEL_CHAIN, models) - .setMissingPredictionTreatment(missingPredictionTreatment); - MiningModel miningModel = new MiningModel(miningFunction, createMiningSchema(models)) - .setSegmentation(segmentation); + .setSegmentation(createSegmentation(Segmentation.MultipleModelMethod.MODEL_CHAIN, missingPredictionTreatment, models)); return miningModel; } @@ -211,11 +208,8 @@ public MiningModel createMultiModelChain(List models, Segmentat } } - Segmentation segmentation = createSegmentation(Segmentation.MultipleModelMethod.MULTI_MODEL_CHAIN, models) - .setMissingPredictionTreatment(missingPredictionTreatment); - MiningModel miningModel = new MiningModel(miningFunction, createMiningSchema(models)) - .setSegmentation(segmentation); + .setSegmentation(createSegmentation(Segmentation.MultipleModelMethod.MULTI_MODEL_CHAIN, missingPredictionTreatment, models)); return miningModel; } @@ -248,12 +242,12 @@ public MiningSchema createMiningSchema(List models){ } static - public Segmentation createSegmentation(Segmentation.MultipleModelMethod multipleModelMethod, List models){ - return createSegmentation(multipleModelMethod, models, null); + public Segmentation createSegmentation(Segmentation.MultipleModelMethod multipleModelMethod, Segmentation.MissingPredictionTreatment missingPredictionTreatment, List models){ + return createSegmentation(multipleModelMethod, missingPredictionTreatment, models, null); } static - public Segmentation createSegmentation(Segmentation.MultipleModelMethod multipleModelMethod, List models, List weights){ + public Segmentation createSegmentation(Segmentation.MultipleModelMethod multipleModelMethod, Segmentation.MissingPredictionTreatment missingPredictionTreatment, List models, List weights){ if((weights != null) && (models.size() != weights.size())){ throw new IllegalArgumentException(); @@ -275,7 +269,10 @@ public Segmentation createSegmentation(Segmentation.MultipleModelMethod multiple segments.add(segment); } - return new Segmentation(multipleModelMethod, segments); + Segmentation segmentation = new Segmentation(multipleModelMethod, segments) + .setMissingPredictionTreatment(missingPredictionTreatment); + + return segmentation; } static