WIP

Author: FlorentinD

graphdatascience/gnn/gnn_nc_runner.py

@@ -1,20 +1,13 @@
-import json
 from typing import Any, List
 
 from ..error.illegal_attr_checker import IllegalAttrChecker
 from ..error.uncallable_namespace import UncallableNamespace
+import json
 
 
 class GNNNodeClassificationRunner(UncallableNamespace, IllegalAttrChecker):
-    def train(
-        self,
-        graph_name: str,
-        model_name: str,
-        feature_properties: List[str],
-        target_property: str,
-        target_node_label: str = None,
-        node_labels: List[str] = None,
-    ) -> "Series[Any]":
+    def train(self, graph_name: str, model_name: str, feature_properties: List[str], target_property: str,
+              target_node_label: str = None, node_labels: List[str] = None) -> "Series[Any]":
         configMap = {
             "featureProperties": feature_properties,
             "targetProperty": target_property,
@@ -27,18 +20,21 @@ def train(
         mlTrainingConfig = json.dumps(configMap)
         # TODO query available node labels
         node_labels = ["Paper"] if not node_labels else node_labels
+
+        # use arrow direclty here
         self._query_runner.run_query(
-            f"CALL gds.upload.graph('{graph_name}', {{mlTrainingConfig: '{mlTrainingConfig}', modelName: '{model_name}', nodeLabels: {node_labels}, nodeProperties: {node_properties}}})"
-        )
+            "CALL gds.upload.graph($config)", 
+            params={"config": 
+                    {"graph_name": graph_name, 
+                     "mlTrainingConfig": mlTrainingConfig, 
+                     "modelName": model_name, 
+                     "nodeLabels": node_labels,
+                    "nodeProperties": node_properties
+                    }
+                    })
+
 
-    def predict(
-        self,
-        graph_name: str,
-        model_name: str,
-        feature_properties: List[str],
-        target_node_label: str = None,
-        node_labels: List[str] = None,
-    ) -> "Series[Any]":
+    def predict(self, graph_name: str, model_name: str, feature_properties: List[str], target_node_label: str = None, node_labels: List[str] = None) -> "Series[Any]":
         configMap = {
             "featureProperties": feature_properties,
             "job_type": "predict",
@@ -49,5 +45,4 @@ def predict(
         # TODO query available node labels
         node_labels = ["Paper"] if not node_labels else node_labels
         self._query_runner.run_query(
-            f"CALL gds.upload.graph('{graph_name}', {{mlTrainingConfig: '{mlTrainingConfig}', modelName: '{model_name}', nodeLabels: {node_labels}, nodeProperties: {feature_properties}}})"
-        )
+            f"CALL gds.upload.graph('{graph_name}', {{mlTrainingConfig: '{mlTrainingConfig}', modelName: '{model_name}', nodeLabels: {node_labels}, nodeProperties: {feature_properties}}})")