Make predict call short, add predictedProbabilityProperty

Co-authored-by: Jacob Sznajdman <breakanalysis@gmail.com>

Author: orazve

graphdatascience/gnn/gnn_nc_runner.py

@@ -7,14 +7,14 @@
 
 class GNNNodeClassificationRunner(UncallableNamespace, IllegalAttrChecker):
     def train(
-        self,
-        graph_name: str,
-        model_name: str,
-        feature_properties: List[str],
-        target_property: str,
-        relationship_types: List[str],
-        target_node_label: str = None,
-        node_labels: List[str] = None,
+            self,
+            graph_name: str,
+            model_name: str,
+            feature_properties: List[str],
+            target_property: str,
+            relationship_types: List[str],
+            target_node_label: str = None,
+            node_labels: List[str] = None,
     ) -> "Series[Any]":  # noqa: F821
         mlConfigMap = {
             "featureProperties": feature_properties,
@@ -40,26 +40,18 @@ def train(
         )
 
     def predict(
-        self,
-        graph_name: str,
-        model_name: str,
-        feature_properties: List[str],
-        relationship_types: List[str],
-        mutateProperty: str,
-        target_node_label: str = None,
-        node_labels: List[str] = None,
+            self,
+            graph_name: str,
+            model_name: str,
+            mutateProperty: str,
+            predictedProbabilityProperty: str = None,
     ) -> "Series[Any]":  # noqa: F821
         mlConfigMap = {
-            "featureProperties": feature_properties,
             "job_type": "predict",
-            "nodeProperties": feature_properties,
-            "relationshipTypes": relationship_types,
             "mutateProperty": mutateProperty
         }
-        if target_node_label:
-            mlConfigMap["targetNodeLabel"] = target_node_label
-        if node_labels:
-            mlConfigMap["nodeLabels"] = node_labels
+        if predictedProbabilityProperty:
+            mlConfigMap["predictedProbabilityProperty"] = predictedProbabilityProperty
 
         mlTrainingConfig = json.dumps(mlConfigMap)
         self._query_runner.run_query(