Remove Model dependency from NodeClassificationTrain

It returns a Classifier directly.
Creating the model is the job of the pipeline.

Author: Mats-SX

pipeline/src/main/java/org/neo4j/gds/ml/pipeline/nodePipeline/classification/train/NodeClassificationTrain.java

@@ -21,7 +21,6 @@
 
 import org.jetbrains.annotations.NotNull;
 import org.neo4j.gds.api.Graph;
-import org.neo4j.gds.core.model.Model;
 import org.neo4j.gds.core.utils.TerminationFlag;
 import org.neo4j.gds.core.utils.mem.MemoryEstimation;
 import org.neo4j.gds.core.utils.mem.MemoryEstimations;
@@ -51,7 +50,6 @@
 import org.neo4j.gds.ml.nodeClassification.ClassificationMetricComputer;
 import org.neo4j.gds.ml.nodePropertyPrediction.NodeSplitter;
 import org.neo4j.gds.ml.pipeline.TrainingStatistics;
-import org.neo4j.gds.ml.pipeline.nodePipeline.NodeClassificationPredictPipeline;
 import org.neo4j.gds.ml.pipeline.nodePipeline.NodePropertyPredictionSplitConfig;
 import org.neo4j.gds.ml.pipeline.nodePipeline.classification.NodeClassificationTrainingPipeline;
 import org.neo4j.gds.ml.splitting.FractionSplitter;
@@ -293,10 +291,7 @@ public NodeClassificationTrainResult compute() {
 
         Classifier retrainedModelData = retrainBestModel(nodeSplits.allTrainingExamples(), trainingStatistics);
 
-        return ImmutableNodeClassificationTrainResult.of(
-            createModel(retrainedModelData, trainingStatistics),
-            trainingStatistics
-        );
+        return ImmutableNodeClassificationTrainResult.of(retrainedModelData, trainingStatistics);
     }
 
     private void selectBestModel(List<TrainingExamplesSplit> nodeSplits, TrainingStatistics trainingStatistics) {
@@ -412,29 +407,6 @@ private Classifier retrainBestModel(HugeLongArray trainSet, TrainingStatistics t
         return retrainedClassifier;
     }
 
-    private Model<Classifier.ClassifierData, NodeClassificationPipelineTrainConfig, NodeClassificationPipelineModelInfo> createModel(
-        Classifier classifier,
-        TrainingStatistics trainingStatistics
-    ) {
-
-        var modelInfo = NodeClassificationPipelineModelInfo.builder()
-            .classes(classIdMap.originalIdsList())
-            .bestParameters(trainingStatistics.bestParameters())
-            .metrics(trainingStatistics.metricsForWinningModel())
-            .pipeline(NodeClassificationPredictPipeline.from(pipeline))
-            .build();
-
-        return Model.of(
-            config.username(),
-            config.modelName(),
-            NodeClassificationTrainingPipeline.MODEL_TYPE,
-            graph.schema(),
-            classifier.data(),
-            config,
-            modelInfo
-        );
-    }
-
     private Classifier trainModel(
         HugeLongArray trainSet,
         TrainerConfig trainerConfig,

pipeline/src/main/java/org/neo4j/gds/ml/pipeline/nodePipeline/classification/train/NodeClassificationTrainPipelineExecutor.java

@@ -31,6 +31,7 @@
 import org.neo4j.gds.ml.pipeline.ImmutableGraphFilter;
 import org.neo4j.gds.ml.pipeline.PipelineExecutor;
 import org.neo4j.gds.ml.pipeline.TrainingStatistics;
+import org.neo4j.gds.ml.pipeline.nodePipeline.NodeClassificationPredictPipeline;
 import org.neo4j.gds.ml.pipeline.nodePipeline.classification.NodeClassificationTrainingPipeline;
 
 import java.util.List;
@@ -108,7 +109,22 @@ protected NodeClassificationTrainPipelineResult execute(Map<DatasetSplits, Graph
             .create(graph, pipeline, config, progressTracker, terminationFlag)
             .compute();
 
-        return ImmutableNodeClassificationTrainPipelineResult.of(trainResult.model(), trainResult.trainingStatistics());
+        var catalogModel = Model.of(
+            config.username(),
+            config.modelName(),
+            NodeClassificationTrainingPipeline.MODEL_TYPE,
+            schemaBeforeSteps,
+            trainResult.classifier().data(),
+            config,
+            NodeClassificationPipelineModelInfo.builder()
+                .classes(trainResult.classifier().classIdMap().originalIdsList())
+                .bestParameters(trainResult.trainingStatistics().bestParameters())
+                .metrics(trainResult.trainingStatistics().metricsForWinningModel())
+                .pipeline(NodeClassificationPredictPipeline.from(pipeline))
+                .build()
+        );
+
+        return ImmutableNodeClassificationTrainPipelineResult.of(catalogModel, trainResult.trainingStatistics());
     }
 
     @ValueClass

pipeline/src/main/java/org/neo4j/gds/ml/pipeline/nodePipeline/classification/train/NodeClassificationTrainResult.java

@@ -20,12 +20,11 @@
 package org.neo4j.gds.ml.pipeline.nodePipeline.classification.train;
 
 import org.neo4j.gds.annotation.ValueClass;
-import org.neo4j.gds.core.model.Model;
 import org.neo4j.gds.ml.models.Classifier;
 import org.neo4j.gds.ml.pipeline.TrainingStatistics;
 
 @ValueClass
 public interface NodeClassificationTrainResult {
-    Model<Classifier.ClassifierData, NodeClassificationPipelineTrainConfig, NodeClassificationPipelineModelInfo> model();
+    Classifier classifier();
     TrainingStatistics trainingStatistics();
 }

pipeline/src/test/java/org/neo4j/gds/ml/pipeline/nodePipeline/classification/train/NodeClassificationTrainTest.java

@@ -149,9 +149,7 @@ void selectsTheBestModel(ClassificationMetricSpecification metricSpecification)
         );
 
         var result = ncTrain.compute();
-        var model = result.model();
 
-        var customInfo = model.customInfo();
         List<ModelStats> validationScores = result.trainingStatistics().getValidationStats(metric);
 
         assertThat(validationScores).hasSize(MAX_TRIALS);
@@ -162,7 +160,7 @@ void selectsTheBestModel(ClassificationMetricSpecification metricSpecification)
                 .isNotCloseTo(validationScores.get(i).avg(), Percentage.withPercentage(0.2));
         }
 
-        var actualWinnerParams = customInfo.bestParameters();
+        var actualWinnerParams = result.trainingStatistics().bestParameters();
         assertThat(actualWinnerParams.toMap()).isEqualTo(expectedWinner.toMap());
     }
 
@@ -226,30 +224,27 @@ void shouldProduceDifferentMetricsForDifferentTrainings(ClassificationMetricSpec
         );
 
         var bananasModelTrainResult = bananasTrain.compute();
-        var bananasModel = bananasModelTrainResult.model();
+        var bananasClassifier = bananasModelTrainResult.classifier();
         var arrayModelTrainResult = arrayPropertyTrain.compute();
-        var arrayPropertyModel = arrayModelTrainResult.model();
+        var arrayPropertyClassifier = arrayModelTrainResult.classifier();
 
-        assertThat(arrayPropertyModel)
+        assertThat(arrayPropertyClassifier)
             .usingRecursiveComparison()
-            .withFailMessage("The trained models are exactly the same instance!")
-            .isNotSameAs(bananasModel);
+            .withFailMessage("The trained classifiers are exactly the same instance!")
+            .isNotSameAs(bananasClassifier);
 
-        assertThat(arrayPropertyModel.data())
+        assertThat(arrayPropertyClassifier.data())
             .usingRecursiveComparison()
             .withFailMessage("Should not produce the same trained `data`!")
-            .isNotEqualTo(bananasModel.data());
+            .isNotEqualTo(bananasClassifier.data());
 
-        var bananasCustomInfo = bananasModel.customInfo();
-        var bananasValidationScore = bananasCustomInfo.metrics().get(metric);
+        var bananasMetrics = bananasModelTrainResult.trainingStatistics().metricsForWinningModel().get(metric);
+        var arrayPropertyMetrics = arrayModelTrainResult.trainingStatistics().metricsForWinningModel().get(metric);
 
-        var arrayPropertyCustomInfo = arrayPropertyModel.customInfo();
-        var arrayPropertyValidationScores = arrayPropertyCustomInfo.metrics().get(metric);
-
-        assertThat(arrayPropertyValidationScores)
+        assertThat(arrayPropertyMetrics)
             .usingRecursiveComparison()
-            .isNotSameAs(bananasValidationScore)
-            .isNotEqualTo(bananasValidationScore);
+            .isNotSameAs(bananasMetrics)
+            .isNotEqualTo(bananasMetrics);
     }
 
     @Test
@@ -450,9 +445,9 @@ void seededNodeClassification(int concurrency) {
         var firstResult = algoSupplier.get().compute();
         var secondResult = algoSupplier.get().compute();
 
-        assertThat(((LogisticRegressionData) firstResult.model().data()).weights().data())
+        assertThat(((LogisticRegressionData) firstResult.classifier().data()).weights().data())
             .matches(matrix -> matrix.equals(
-                ((LogisticRegressionData) secondResult.model().data()).weights().data(),
+                ((LogisticRegressionData) secondResult.classifier().data()).weights().data(),
                 1e-10
             ));
     }