WIP

Author: breakanalysis

graphdatascience/gnn/__init__.py

@@ -0,0 +1,17 @@
+from .gnn_nc_runner import GNNNodeClassificationRunner
+from ..caller_base import CallerBase
+from ..error.illegal_attr_checker import IllegalAttrChecker
+from ..error.uncallable_namespace import UncallableNamespace
+
+class GNNRunner(UncallableNamespace, IllegalAttrChecker):
+    @property
+    def nodeClassification(self) -> GNNNodeClassificationRunner:
+        return GNNNodeClassificationRunner(self._query_runner, f"{self._namespace}.nodeClassification", self._server_version)
+
+class GnnEndpoints(CallerBase):
+    @property
+    def gnn(self) -> GNNRunner:
+        return GNNRunner(self._query_runner, f"{self._namespace}.gnn", self._server_version)
+
+
+

graphdatascience/gnn/gnn_endpoints.py

@@ -0,0 +1,20 @@
+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 = {
+            "feature_properties": feature_properties,
+            "target_property": target_property,
+        }
+        if target_node_label:
+            configMap["targetNodeLabel"] = target_node_label
+        mlTrainingConfig = json.dumps(configMap)
+        # TODO query avaiable 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}}})")

graphdatascience/gnn/gnn_nc_runner.py

@@ -47,6 +47,7 @@
     "gds.alpha.pipeline.nodeRegression.predict.stream",
     "gds.alpha.pipeline.nodeRegression.selectFeatures",
     "gds.alpha.pipeline.nodeRegression.train",
+    "gds.gnn.nc",
     "gds.similarity.cosine",
     "gds.similarity.euclidean",
     "gds.similarity.euclideanDistance",