Finish modifying LP train tasks

Co-Authored-By: Adam Schill Collberg<adam.schill.collberg@protonmail.com>

Author: breakanalysis

pipeline/src/main/java/org/neo4j/gds/ml/pipeline/linkPipeline/train/LinkPredictionTrain.java

@@ -100,25 +100,20 @@ public static List<Task> progressTasks(
         LinkPredictionSplitConfig splitConfig,
         int numberOfModelSelectionTrials
     ) {
-        // the relationship count estimates depend on both UndirectedEdgeSplitter
-        // and the volume set in extractFeaturesAndLabels
-        var selectionRatio = (1 + splitConfig.negativeSamplingRatio());
-        long nonTestRelationshipCount = (long) (relationshipCount * (1 - splitConfig.testFraction()));
-        long testRelationshipCount = (long) (relationshipCount * splitConfig.testFraction() * selectionRatio);
-        long trainRelationshipCount = (long) (nonTestRelationshipCount * splitConfig.trainFraction() * selectionRatio);
+        var sizes = splitConfig.expectedSetSizes(relationshipCount);
         return List.of(
-            Tasks.leaf("Extract train features", trainRelationshipCount),
+            Tasks.leaf("Extract train features", sizes.trainSize()),
             Tasks.iterativeFixed(
                 "Select best model",
-                () -> List.of(Tasks.leaf("Trial", splitConfig.validationFolds() * trainRelationshipCount)),
+                () -> List.of(Tasks.leaf("Trial", splitConfig.validationFolds() * sizes.trainSize())),
                 numberOfModelSelectionTrials
             ),
-            ClassifierTrainer.progressTask("Train best model", trainRelationshipCount),
-            Tasks.leaf("Compute train metrics"),
+            ClassifierTrainer.progressTask("Train best model", sizes.trainSize()),
+            Tasks.leaf("Compute train metrics", sizes.trainSize()),
             Tasks.task(
                 "Evaluate on test data",
-                Tasks.leaf("Extract test features", testRelationshipCount),
-                Tasks.leaf("Compute test metrics")
+                Tasks.leaf("Extract test features", sizes.testSize()),
+                Tasks.leaf("Compute test metrics", sizes.testSize())
             )
         );
     }

pipeline/src/test/java/org/neo4j/gds/ml/pipeline/linkPipeline/train/LinkPredictionTrainTest.java

@@ -616,7 +616,7 @@ void logProgressLRWithRange() {
         var log = Neo4jProxy.testLog();
         var progressTracker = new TestProgressTracker(
             progressTask(
-                trainGraph.relationshipCount(),
+                2 * trainGraph.relationshipCount(),
                 pipeline.splitConfig(),
                 pipeline.numberOfModelSelectionTrials()
             ),