TensorflowGNN transformer

.github/workflows/build-and-test.yml

@@ -55,8 +55,10 @@ jobs:
           poetry install --all-extras
           poe install-pyg-cpu
           poe install-dgl
+          poe install-tfgnn
       - name: Run Tests
         run: |
+          export TF_USE_LEGACY_KERAS="1"
           poetry run pytest -vvv -m "not slow and not ubuntu and not docker"
       - name: Use the Upload Artifact GitHub Action
         uses: actions/upload-artifact@v4
@@ -131,6 +133,8 @@ jobs:
         run: |
           poetry install --all-extras
           poe install-pyg-cpu
+          poe install-tfgnn
+          export TF_USE_LEGACY_KERAS="1"
           poetry run pytest -vvv -m "not slow and not ubuntu and not docker"
       - name: Save Memgraph Logs
         uses: actions/upload-artifact@v4

README.md

@@ -84,6 +84,7 @@ poetry install # No extras
 poetry install -E arrow # Support for the CSV, Parquet, ORC and IPC/Feather/Arrow formats
 poetry install -E dgl # DGL support (also includes torch)
 poetry install -E docker # Docker support
+poetry install -E tfgnn # TFGNN support
 ```
 
 To run the tests, make sure you have an [active Memgraph instance](https://memgraph.com/docs/getting-started), and execute one of the following commands:
@@ -100,6 +101,7 @@ If you’ve installed only certain extras, it’s also possible to run their ass
 poetry run pytest . -k "arrow"
 poetry run pytest . -k "dgl"
 poetry run pytest . -k "docker"
+poetry run pytest . -k "tfgnn"
 ```
 
 ## Development (how to build)

gqlalchemy/transformations/constants.py

@@ -9,3 +9,4 @@
 NUM_NODES = "num_nodes"
 DEFAULT_NODE_LABEL = "NODE"
 DEFAULT_EDGE_TYPE = "RELATIONSHIP"
+TFGNN_ID = "tfgnn_id"

gqlalchemy/transformations/export/graph_transporter.py

@@ -29,6 +29,11 @@
 except ModuleNotFoundError:
     PyGTranslator = None
 
+try:
+    from gqlalchemy.transformations.translators.tfgnn_translator import TFGNNTranslator
+except ModuleNotFoundError:
+    TFGNNTranslator = None
+
 
 class GraphTransporter(Transporter):
     """Here is a possible example for using this module:
@@ -52,7 +57,7 @@ def __init__(
         >>> transporter = GraphTransporter("dgl")
         graph = transporter.export()
         Args:
-            graph_type: dgl, pyg or nx
+            graph_type: dgl, pyg, nx or tfgnn
         """
         super().__init__()
         self.graph_type = graph_type.upper()
@@ -64,8 +69,11 @@ def __init__(
             self.translator = PyGTranslator(host, port, username, password, encrypted, client_name, lazy)
         elif self.graph_type == GraphType.NX.name:
             self.translator = NxTranslator(host, port, username, password, encrypted, client_name, lazy)
+        elif self.graph_type == GraphType.TFGNN.name:
+            raise_if_not_imported(dependency=TFGNNTranslator, dependency_name="tensorflow-gnn")
+            self.translator = TFGNNTranslator(host, port, username, password, encrypted, client_name, lazy)
         else:
-            raise ValueError("Unknown export option. Currently supported are DGL, PyG and NetworkX.")
+            raise ValueError("Unknown export option. Currently supported are DGL, PyG, NetworkX and TFGNN.")
 
     def export(self):
         """Creates graph instance for the wanted export option."""

gqlalchemy/transformations/graph_type.py

@@ -1,3 +1,3 @@
 from enum import Enum
 
-GraphType = Enum("GraphType", ["DGL", "PYG", "NX"])
+GraphType = Enum("GraphType", ["DGL", "PYG", "NX", "TFGNN"])