Cleanup nc_runner

Co-authored-by: Florentin Dörre <[email protected]>
Co-authored-by: Olga Razvenskaia <[email protected]>

Author: 3 people

graphdatascience/gnn/gnn_nc_runner.py

@@ -1,46 +1,66 @@
+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]":
-        configMap = {
+    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]":  # noqa: F821
+        mlConfigMap = {
             "featureProperties": feature_properties,
             "targetProperty": target_property,
             "job_type": "train",
-            "nodeProperties": feature_properties + [target_property]
+            "nodeProperties": feature_properties + [target_property],
         }
 
         if target_node_label:
-            configMap["targetNodeLabel"] = target_node_label
+            mlConfigMap["targetNodeLabel"] = target_node_label
         if node_labels:
-            configMap["nodeLabels"] = node_labels
+            mlConfigMap["nodeLabels"] = node_labels
 
-        mlTrainingConfig = json.dumps(configMap)
+        mlTrainingConfig = json.dumps(mlConfigMap)
 
         # token and uri will be injected by arrow_query_runner
         self._query_runner.run_query(
-            f"CALL gds.upload.graph($graph_name, $config)", 
-            params={"graph_name": graph_name, "config": {
-                     "mlTrainingConfig": mlTrainingConfig, 
-                     "modelName": model_name
-                    }}
-                    )
-
+            "CALL gds.upload.graph($graph_name, $config)",
+            params={
+                "graph_name": graph_name,
+                "config": {"mlTrainingConfig": mlTrainingConfig, "modelName": model_name},
+            },
+        )
 
-    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 = {
+    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]":  # noqa: F821
+        mlConfigMap = {
             "featureProperties": feature_properties,
             "job_type": "predict",
+            "nodeProperties": feature_properties,
         }
         if target_node_label:
-            configMap["targetNodeLabel"] = target_node_label
-        mlTrainingConfig = json.dumps(configMap)
-        # TODO query available node labels
-        node_labels = ["Paper"] if not node_labels else node_labels
+            mlConfigMap["targetNodeLabel"] = target_node_label
+        if node_labels:
+            mlConfigMap["nodeLabels"] = node_labels
+
+        mlTrainingConfig = json.dumps(mlConfigMap)
         self._query_runner.run_query(
-            f"CALL gds.upload.graph('{graph_name}', {{mlTrainingConfig: '{mlTrainingConfig}', modelName: '{model_name}', nodeLabels: {node_labels}, nodeProperties: {feature_properties}}})")
+            "CALL gds.upload.graph($graph_name, $config)",
+            params={
+                "graph_name": graph_name,
+                "config": {"mlTrainingConfig": mlTrainingConfig, "modelName": model_name},
+            },
+        )  # type: ignore

graphdatascience/query_runner/arrow_query_runner.py

@@ -191,7 +191,7 @@ def create_graph_constructor(
         return ArrowGraphConstructor(
             database, graph_name, self._flight_client, concurrency, undirected_relationship_types
         )
-    
+
     def _get_or_request_token(self) -> str:
         self._flight_client.authenticate_basic_token(self._auth[0], self._auth[1])
         return self._auth_factory.token()
@@ -232,7 +232,7 @@ def received_headers(self, headers: Dict[str, Any]) -> None:
         auth_header: str = headers.get("authorization", None)
         if not auth_header:
             return
-        # authenticate_basic_token() returns a list. 
+        # authenticate_basic_token() returns a list.
         # TODO We should take the first Bearer element here
         if isinstance(auth_header, list):
             auth_header = auth_header[0]