Training and inference.

examples/Model definition

@@ -0,0 +1,334 @@
+{
+  "edges": [
+    {
+      "id": "Input: embedding 1 Linear 1",
+      "source": "Input: embedding 1",
+      "sourceHandle": "x",
+      "target": "Linear 1",
+      "targetHandle": "x"
+    },
+    {
+      "id": "Input: label 1 MSE loss 1",
+      "source": "Input: label 1",
+      "sourceHandle": "y",
+      "target": "MSE loss 1",
+      "targetHandle": "y"
+    },
+    {
+      "id": "Linear 1 Activation 2",
+      "source": "Linear 1",
+      "sourceHandle": "x",
+      "target": "Activation 2",
+      "targetHandle": "x"
+    },
+    {
+      "id": "Activation 2 MSE loss 1",
+      "source": "Activation 2",
+      "sourceHandle": "x",
+      "target": "MSE loss 1",
+      "targetHandle": "x"
+    },
+    {
+      "id": "MSE loss 1 Optimizer 2",
+      "source": "MSE loss 1",
+      "sourceHandle": "loss",
+      "target": "Optimizer 2",
+      "targetHandle": "loss"
+    }
+  ],
+  "env": "PyTorch model",
+  "nodes": [
+    {
+      "data": {
+        "display": null,
+        "error": null,
+        "meta": {
+          "inputs": {},
+          "name": "Input: embedding",
+          "outputs": {
+            "x": {
+              "name": "x",
+              "position": "top",
+              "type": {
+                "type": "tensor"
+              }
+            }
+          },
+          "params": {},
+          "type": "basic"
+        },
+        "params": {},
+        "status": "planned",
+        "title": "Input: embedding"
+      },
+      "dragHandle": ".bg-primary",
+      "height": 200.0,
+      "id": "Input: embedding 1",
+      "position": {
+        "x": 91.0,
+        "y": 266.0
+      },
+      "type": "basic",
+      "width": 200.0
+    },
+    {
+      "data": {
+        "display": null,
+        "error": null,
+        "meta": {
+          "inputs": {
+            "x": {
+              "name": "x",
+              "position": "bottom",
+              "type": {
+                "type": "tensor"
+              }
+            }
+          },
+          "name": "Linear",
+          "outputs": {
+            "x": {
+              "name": "x",
+              "position": "top",
+              "type": {
+                "type": "tensor"
+              }
+            }
+          },
+          "params": {
+            "output_dim": {
+              "default": "same",
+              "name": "output_dim",
+              "type": {
+                "type": "<class 'str'>"
+              }
+            }
+          },
+          "type": "basic"
+        },
+        "params": {
+          "output_dim": "same"
+        },
+        "status": "planned",
+        "title": "Linear"
+      },
+      "dragHandle": ".bg-primary",
+      "height": 200.0,
+      "id": "Linear 1",
+      "position": {
+        "x": 86.0,
+        "y": 33.0
+      },
+      "type": "basic",
+      "width": 200.0
+    },
+    {
+      "data": {
+        "display": null,
+        "error": null,
+        "meta": {
+          "inputs": {
+            "x": {
+              "name": "x",
+              "position": "bottom",
+              "type": {
+                "type": "tensor"
+              }
+            },
+            "y": {
+              "name": "y",
+              "position": "bottom",
+              "type": {
+                "type": "tensor"
+              }
+            }
+          },
+          "name": "MSE loss",
+          "outputs": {
+            "loss": {
+              "name": "loss",
+              "position": "top",
+              "type": {
+                "type": "tensor"
+              }
+            }
+          },
+          "params": {},
+          "type": "basic"
+        },
+        "params": {},
+        "status": "planned",
+        "title": "MSE loss"
+      },
+      "dragHandle": ".bg-primary",
+      "height": 200.0,
+      "id": "MSE loss 1",
+      "position": {
+        "x": 315.0,
+        "y": -510.0
+      },
+      "type": "basic",
+      "width": 200.0
+    },
+    {
+      "data": {
+        "display": null,
+        "error": null,
+        "meta": {
+          "inputs": {},
+          "name": "Input: label",
+          "outputs": {
+            "y": {
+              "name": "y",
+              "position": "top",
+              "type": {
+                "type": "tensor"
+              }
+            }
+          },
+          "params": {},
+          "type": "basic"
+        },
+        "params": {},
+        "status": "planned",
+        "title": "Input: label"
+      },
+      "dragHandle": ".bg-primary",
+      "height": 200.0,
+      "id": "Input: label 1",
+      "position": {
+        "x": 615.0,
+        "y": -165.0
+      },
+      "type": "basic",
+      "width": 200.0
+    },
+    {
+      "data": {
+        "__execution_delay": 0.0,
+        "collapsed": null,
+        "display": null,
+        "error": null,
+        "meta": {
+          "inputs": {
+            "x": {
+              "name": "x",
+              "position": "bottom",
+              "type": {
+                "type": "tensor"
+              }
+            }
+          },
+          "name": "Activation",
+          "outputs": {
+            "x": {
+              "name": "x",
+              "position": "top",
+              "type": {
+                "type": "tensor"
+              }
+            }
+          },
+          "params": {
+            "type": {
+              "default": "ReLU",
+              "name": "type",
+              "type": {
+                "enum": [
+                  "ReLU",
+                  "Leaky ReLU",
+                  "Tanh",
+                  "Mish"
+                ]
+              }
+            }
+          },
+          "position": {
+            "x": 419.0,
+            "y": 396.0
+          },
+          "type": "basic"
+        },
+        "params": {
+          "type": "Leaky ReLU"
+        },
+        "status": "planned",
+        "title": "Activation"
+      },
+      "dragHandle": ".bg-primary",
+      "height": 200.0,
+      "id": "Activation 2",
+      "position": {
+        "x": 93.61643829835265,
+        "y": -229.04087132886406
+      },
+      "type": "basic",
+      "width": 200.0
+    },
+    {
+      "data": {
+        "__execution_delay": 0.0,
+        "collapsed": null,
+        "display": null,
+        "error": null,
+        "meta": {
+          "inputs": {
+            "loss": {
+              "name": "loss",
+              "position": "bottom",
+              "type": {
+                "type": "tensor"
+              }
+            }
+          },
+          "name": "Optimizer",
+          "outputs": {},
+          "params": {
+            "lr": {
+              "default": 0.001,
+              "name": "lr",
+              "type": {
+                "type": "<class 'float'>"
+              }
+            },
+            "type": {
+              "default": "AdamW",
+              "name": "type",
+              "type": {
+                "enum": [
+                  "AdamW",
+                  "Adafactor",
+                  "Adagrad",
+                  "SGD",
+                  "Lion",
+                  "Paged AdamW",
+                  "Galore AdamW"
+                ]
+              }
+            }
+          },
+          "position": {
+            "x": 526.0,
+            "y": 116.0
+          },
+          "type": "basic"
+        },
+        "params": {
+          "lr": "0.1",
+          "type": "SGD"
+        },
+        "status": "planned",
+        "title": "Optimizer"
+      },
+      "dragHandle": ".bg-primary",
+      "height": 200.0,
+      "id": "Optimizer 2",
+      "position": {
+        "x": 305.6132943499785,
+        "y": -804.0094318451224
+      },
+      "type": "basic",
+      "width": 200.0
+    }
+  ]
+}

examples/Model use

@@ -0,0 +1,288 @@
+{
+  "edges": [
+    {
+      "id": "Import Parquet 1 Train/test split 1",
+      "source": "Import Parquet 1",
+      "sourceHandle": "output",
+      "target": "Train/test split 1",
+      "targetHandle": "bundle"
+    },
+    {
+      "id": "Train/test split 1 Train model 3",
+      "source": "Train/test split 1",
+      "sourceHandle": "output",
+      "target": "Train model 3",
+      "targetHandle": "bundle"
+    },
+    {
+      "id": "Train model 3 Model inference 1",
+      "source": "Train model 3",
+      "sourceHandle": "output",
+      "target": "Model inference 1",
+      "targetHandle": "bundle"
+    }
+  ],
+  "env": "LynxKite Graph Analytics",
+  "nodes": [
+    {
+      "data": {
+        "__execution_delay": 0.0,
+        "collapsed": null,
+        "display": null,
+        "error": null,
+        "meta": {
+          "inputs": {
+            "bundle": {
+              "name": "bundle",
+              "position": "left",
+              "type": {
+                "type": "<class 'lynxkite_graph_analytics.core.Bundle'>"
+              }
+            }
+          },
+          "name": "Train/test split",
+          "outputs": {
+            "output": {
+              "name": "output",
+              "position": "right",
+              "type": {
+                "type": "None"
+              }
+            }
+          },
+          "params": {
+            "table_name": {
+              "default": null,
+              "name": "table_name",
+              "type": {
+                "type": "<class 'str'>"
+              }
+            },
+            "test_ratio": {
+              "default": 0.1,
+              "name": "test_ratio",
+              "type": {
+                "type": "<class 'float'>"
+              }
+            }
+          },
+          "type": "basic"
+        },
+        "params": {
+          "table_name": "df",
+          "test_ratio": 0.1
+        },
+        "status": "done",
+        "title": "Train/test split"
+      },
+      "dragHandle": ".bg-primary",
+      "height": 282.0,
+      "id": "Train/test split 1",
+      "position": {
+        "x": 345.0,
+        "y": 139.0
+      },
+      "type": "basic",
+      "width": 259.0
+    },
+    {
+      "data": {
+        "__execution_delay": 0.0,
+        "collapsed": null,
+        "display": null,
+        "error": null,
+        "meta": {
+          "inputs": {},
+          "name": "Import Parquet",
+          "outputs": {
+            "output": {
+              "name": "output",
+              "position": "right",
+              "type": {
+                "type": "None"
+              }
+            }
+          },
+          "params": {
+            "filename": {
+              "default": null,
+              "name": "filename",
+              "type": {
+                "type": "<class 'str'>"
+              }
+            }
+          },
+          "type": "basic"
+        },
+        "params": {
+          "filename": "uploads/plus-one-dataset.parquet"
+        },
+        "status": "done",
+        "title": "Import Parquet"
+      },
+      "dragHandle": ".bg-primary",
+      "height": 403.0,
+      "id": "Import Parquet 1",
+      "position": {
+        "x": -166.0,
+        "y": 112.0
+      },
+      "type": "basic",
+      "width": 371.0
+    },
+    {
+      "data": {
+        "__execution_delay": 0.0,
+        "collapsed": null,
+        "display": null,
+        "error": null,
+        "meta": {
+          "inputs": {
+            "bundle": {
+              "name": "bundle",
+              "position": "left",
+              "type": {
+                "type": "<class 'lynxkite_graph_analytics.core.Bundle'>"
+              }
+            }
+          },
+          "name": "Train model",
+          "outputs": {
+            "output": {
+              "name": "output",
+              "position": "right",
+              "type": {
+                "type": "None"
+              }
+            }
+          },
+          "params": {
+            "epochs": {
+              "default": 1.0,
+              "name": "epochs",
+              "type": {
+                "type": "<class 'int'>"
+              }
+            },
+            "input_mapping": {
+              "default": null,
+              "name": "input_mapping",
+              "type": {
+                "type": "<class 'str'>"
+              }
+            },
+            "model_workspace": {
+              "default": null,
+              "name": "model_workspace",
+              "type": {
+                "type": "<class 'str'>"
+              }
+            },
+            "save_as": {
+              "default": "model",
+              "name": "save_as",
+              "type": {
+                "type": "<class 'str'>"
+              }
+            }
+          },
+          "position": {
+            "x": 675.0,
+            "y": 144.0
+          },
+          "type": "basic"
+        },
+        "params": {
+          "epochs": "1000",
+          "input_mapping": "{\"Input__embedding_1_x\": {\"df\": \"df_train\", \"column\": \"x\"}, \"Input__label_1_y\": {\"df\": \"df_train\", \"column\": \"y\"  }}",
+          "model_workspace": "Model definition",
+          "save_as": "model"
+        },
+        "status": "done",
+        "title": "Train model"
+      },
+      "dragHandle": ".bg-primary",
+      "height": 519.0,
+      "id": "Train model 3",
+      "position": {
+        "x": 687.3818749999999,
+        "y": -34.16902777777775
+      },
+      "type": "basic",
+      "width": 640.0
+    },
+    {
+      "data": {
+        "__execution_delay": 0.0,
+        "collapsed": null,
+        "display": null,
+        "error": null,
+        "meta": {
+          "inputs": {
+            "bundle": {
+              "name": "bundle",
+              "position": "left",
+              "type": {
+                "type": "<class 'lynxkite_graph_analytics.core.Bundle'>"
+              }
+            }
+          },
+          "name": "Model inference",
+          "outputs": {
+            "output": {
+              "name": "output",
+              "position": "right",
+              "type": {
+                "type": "None"
+              }
+            }
+          },
+          "params": {
+            "input_mapping": {
+              "default": "",
+              "name": "input_mapping",
+              "type": {
+                "type": "<class 'str'>"
+              }
+            },
+            "model_name": {
+              "default": "model",
+              "name": "model_name",
+              "type": {
+                "type": "<class 'str'>"
+              }
+            },
+            "output_mapping": {
+              "default": "",
+              "name": "output_mapping",
+              "type": {
+                "type": "<class 'str'>"
+              }
+            }
+          },
+          "position": {
+            "x": 506.0,
+            "y": 115.0
+          },
+          "type": "basic"
+        },
+        "params": {
+          "input_mapping": "{\"Input__embedding_1_x\": {\"df\": \"df_test\", \"column\": \"x\"}}",
+          "model_name": "model",
+          "output_mapping": "{\"Activation_2_x\": {\"df\": \"df_test\", \"column\": \"predicted\"}}"
+        },
+        "status": "done",
+        "title": "Model inference"
+      },
+      "dragHandle": ".bg-primary",
+      "height": 429.0,
+      "id": "Model inference 1",
+      "position": {
+        "x": 1445.5664910683593,
+        "y": 12.075943590382515
+      },
+      "type": "basic",
+      "width": 410.0
+    }
+  ]
+}

lynxkite-core/src/lynxkite/core/ops.py

@@ -61,7 +61,7 @@ class Parameter(BaseConfig):
     @staticmethod
     def options(name, options, default=None):
         e = enum.Enum(f"OptionsFor_{name}", options)
-        return Parameter.basic(name, e[default or options[0]], e)
+        return Parameter.basic(name, default or options[0], e)
 
     @staticmethod
     def collapsed(name, default, type=None):
@@ -154,9 +154,7 @@ class Op(BaseConfig):
 
     def __call__(self, *inputs, **params):
         # Convert parameters.
-        for p in params:
-            if p in self.params:
-                params[p] = _param_to_type(p, params[p], self.params[p].type)
+        params = self.convert_params(params)
         res = self.func(*inputs, **params)
         if not isinstance(res, Result):
             # Automatically wrap the result in a Result object, if it isn't already.
@@ -172,6 +170,15 @@ class Op(BaseConfig):
                 res.display = res.output
         return res
 
+    def convert_params(self, params):
+        """Returns the parameters converted to the expected type."""
+        res = {}
+        for p in params:
+            res[p] = params[p]
+            if p in self.params:
+                res[p] = _param_to_type(p, params[p], self.params[p].type)
+        return res
+
 
 def op(env: str, name: str, *, view="basic", outputs=None, params=None):
     """Decorator for defining an operation."""

lynxkite-graph-analytics/src/lynxkite_graph_analytics/core.py

@@ -42,7 +42,7 @@ class Bundle:
 
     dfs: dict[str, pd.DataFrame] = dataclasses.field(default_factory=dict)
     relations: list[RelationDefinition] = dataclasses.field(default_factory=list)
-    other: dict[str, typing.Any] = None
+    other: dict[str, typing.Any] = dataclasses.field(default_factory=dict)
 
     @classmethod
     def from_nx(cls, graph: nx.Graph):
@@ -102,7 +102,7 @@ class Bundle:
         return Bundle(
             dfs=dict(self.dfs),
             relations=list(self.relations),
-            other=dict(self.other) if self.other else None,
+            other=dict(self.other),
         )
 
     def to_dict(self, limit: int = 100):

lynxkite-graph-analytics/src/lynxkite_graph_analytics/lynxkite_ops.py

@@ -2,10 +2,14 @@
 
 import enum
 import os
+import pathlib
 import fsspec
 from lynxkite.core import ops
 from collections import deque
-from . import core
+
+from tqdm import tqdm
+from . import core, pytorch_model_ops
+from lynxkite.core import workspace
 import grandcypher
 import joblib
 import matplotlib
@@ -344,10 +348,13 @@ def create_graph(bundle: core.Bundle, *, relations: str = None) -> core.Bundle:
     return ops.Result(output=bundle, display=bundle.to_dict(limit=100))
 
 
-@op("Define model")
-def define_model(*, model_workspace: str, save_as: str = "model"):
-    """Reads a PyTorch model workspace and returns it as a model in a bundle."""
-    return None
+def load_ws(model_workspace: str):
+    cwd = pathlib.Path()
+    path = cwd / model_workspace
+    assert path.is_relative_to(cwd)
+    assert path.exists(), f"Workspace {path} does not exist"
+    ws = workspace.load(path)
+    return ws
 
 
 @op("Biomedical foundation graph (PLACEHOLDER)")
@@ -358,25 +365,54 @@ def biomedical_foundation_graph(*, filter_nodes: str):
 
 @op("Train model")
 def train_model(
-    bundle: core.Bundle, *, model_name: str, model_mapping: str, epochs: int = 1
+    bundle: core.Bundle,
+    *,
+    model_workspace: str,
+    input_mapping: str,
+    epochs: int = 1,
+    save_as: str = "model",
 ):
     """Trains the selected model on the selected dataset. Most training parameters are set in the model definition."""
-    return None
+    ws = load_ws(model_workspace)
+    input_mapping = json.loads(input_mapping)
+    inputs = pytorch_model_ops.to_tensors(bundle, input_mapping)
+    m = pytorch_model_ops.build_model(ws, inputs)
+    t = tqdm(range(epochs), desc="Training model")
+    for _ in t:
+        loss = m.train(inputs)
+        t.set_postfix({"loss": loss})
+    bundle = bundle.copy()
+    bundle.other[save_as] = m
+    return bundle
 
 
 @op("Model inference")
 def model_inference(
     bundle: core.Bundle,
     *,
-    model_name: str,
-    model_mapping: str,
-    save_output_as: str = "prediction",
+    model_name: str = "model",
+    input_mapping: str = "",
+    output_mapping: str = "",
 ):
     """Executes a trained model."""
-    return None
+    m = bundle.other[model_name]
+    input_mapping = json.loads(input_mapping)
+    output_mapping = json.loads(output_mapping)
+    inputs = pytorch_model_ops.to_tensors(bundle, input_mapping)
+    outputs = m.inference(inputs)
+    bundle = bundle.copy()
+    for k, v in output_mapping.items():
+        bundle.dfs[v["df"]][v["column"]] = outputs[k].detach().numpy().tolist()
+    return bundle
 
 
 @op("Train/test split")
 def train_test_split(bundle: core.Bundle, *, table_name: str, test_ratio: float = 0.1):
     """Splits a dataframe in the bundle into separate "_train" and "_test" dataframes."""
-    return None
+    df = bundle.dfs[table_name]
+    test = df.sample(frac=test_ratio)
+    train = df.drop(test.index)
+    bundle = bundle.copy()
+    bundle.dfs[f"{table_name}_train"] = train
+    bundle.dfs[f"{table_name}_test"] = test
+    return bundle

lynxkite-graph-analytics/src/lynxkite_graph_analytics/pytorch_model_ops.py

@@ -6,6 +6,7 @@ from lynxkite.core.ops import Parameter as P
 import torch
 import torch_geometric as pyg
 from dataclasses import dataclass
+from . import core
 
 ENV = "PyTorch model"
 
@@ -162,11 +163,18 @@ class ModelConfig:
         self.optimizer.step()
         return loss.item()
 
+    def copy(self):
+        """Returns a copy of the model."""
+        c = super().copy()
+        c.model = self.model.copy()
+        return c
+
 
 def build_model(
     ws: workspace.Workspace, inputs: dict[str, torch.Tensor]
 ) -> ModelConfig:
     """Builds the model described in the workspace."""
+    catalog = ops.CATALOGS[ENV]
     optimizers = []
     nodes = {}
     for node in ws.nodes:
@@ -197,7 +205,8 @@ def build_model(
     for node_id in ts.static_order():
         node = nodes[node_id]
         t = node.data.title
-        p = node.data.params
+        op = catalog[t]
+        p = op.convert_params(node.data.params)
         for b in dependencies[node_id]:
             if b in in_loss:
                 in_loss.add(node_id)
@@ -216,7 +225,9 @@ def build_model(
                 [(ib, ih)] = edges[node_id, "x"]
                 i = _to_id(ib) + "_" + ih
                 used_inputs.add(i)
-                f = getattr(torch.nn.functional, p["type"].lower().replace(" ", "_"))
+                f = getattr(
+                    torch.nn.functional, p["type"].name.lower().replace(" ", "_")
+                )
                 ls.append((f, f"{i} -> {nid}_x"))
                 sizes[f"{nid}_x"] = sizes[i]
             case "MSE loss":
@@ -248,7 +259,18 @@ def build_model(
         f"loss should have no parameters: {list(cfg['loss'].parameters())}"
     )
     # Create optimizer.
-    p = nodes[optimizer].data.params
-    o = getattr(torch.optim, p["type"])
+    op = catalog["Optimizer"]
+    p = op.convert_params(nodes[optimizer].data.params)
+    o = getattr(torch.optim, p["type"].name)
     cfg["optimizer"] = o(cfg["model"].parameters(), lr=p["lr"])
     return ModelConfig(**cfg)
+
+
+def to_tensors(b: core.Bundle, m: dict[str, dict]) -> dict[str, torch.Tensor]:
+    """Converts a tensor to the correct type for PyTorch."""
+    tensors = {}
+    for k, v in m.items():
+        tensors[k] = torch.tensor(
+            b.dfs[v["df"]][v["column"]].to_list(), dtype=torch.float32
+        )
+    return tensors