# GNN

relationalai.semantics.reasoners.predictive.estimators.relational.gnn

```python
GNN(
    *,
    exp_database: str,
    exp_schema: str,
    database: Optional[str] = None,
    schema: Optional[str] = None,
    graph: Optional[Graph] = None,
    property_transformer: Optional[PropertyTransformer] = None,
    train: Optional[b.Relationship | b.Fragment | b.Chain] = None,
    validation: Optional[b.Relationship | b.Fragment | b.Chain] = None,
    source_concept: Optional[b.Concept] = None,
    target_concept: Optional[b.Concept] = None,
    task_type: Optional[str] = None,
    eval_metric: Optional[str] = None,
    use_current_time: bool = True,
    has_time_column: Optional[bool] = None,
    model_database: Optional[str] = None,
    model_schema: Optional[str] = None,
    model_name: Optional[str] = None,
    version_name: Optional[str] = None,
    model_run_id: Optional[str] = None,
    test_batch_size: Optional[int] = None,
    stream_logs: bool = True,
    extract_embeddings: bool = False,
    dataset_alias: Optional[str] = None,
    parallel_reasoners_init: bool = True,
    **train_params: Any
)
```

Train, load, register and predict with a Graph Neural Network.

`GNN` supports two workflows:

- **Fit workflow** — provide `graph`, `train`, `validation`, and a
  `task_type`, then call [`GNN.fit`](/api/python/v1.35/semantics/reasoners/predictive/estimators/relational/gnn/gnn_class/index.md#method-fit) followed by [`GNN.predictions`](/api/python/v1.35/semantics/reasoners/predictive/estimators/relational/gnn/gnn_class/index.md#method-predictions).
- **Load workflow** — provide a previously trained model via
  `model_run_id` **or** the four registry parameters
  (`model_database`, `model_schema`, `model_name`,
  `version_name`), then call [`GNN.load`](/api/python/v1.35/semantics/reasoners/predictive/estimators/relational/gnn/gnn_class/index.md#method-load) followed by
  [`GNN.predictions`](/api/python/v1.35/semantics/reasoners/predictive/estimators/relational/gnn/gnn_class/index.md#method-predictions).

In either workflow you can register the trained model in
the Snowflake Model Registry using the [`GNN.register_model`](/api/python/v1.35/semantics/reasoners/predictive/estimators/relational/gnn/gnn_class/index.md#method-register_model) method.

## Parameters

- `exp_database` (`str`) - Snowflake database for experiment storage.
- `exp_schema` (`str`) - Snowflake schema for experiment storage.
- `database` (`str`, default: `None`) - Snowflake database to save predictions in.
- `schema` (`str`, default: `None`) - Snowflake schema to save predictions in.
- `graph` (`Graph`, default: `None`) - The knowledge graph with edges defined.  Required for the fit workflow.
- `property_transformer` (`PropertyTransformer`, default: `None`) - Column-level semantic type annotations.  If omitted, all column types are auto-inferred.
- `train` (`Relationship or Fragment`, default: `None`) - Training split relationship.  Required for the fit workflow.
- `validation` (`Relationship or Fragment`, default: `None`) - Validation split relationship.  Required for the fit workflow.
- `source_concept` (`Concept`, default: `None`) - Source concept for the load workflow (inferred from `train` in the fit workflow).
- `target_concept` (`Concept`, default: `None`) - Target concept for link-prediction tasks in the load workflow (inferred from `train` in the fit workflow).
- `task_type` (`str`, default: `None`) - One of `"binary_classification"`, `"multiclass_classification"`, `"multilabel_classification"`, `"regression"`, `"link_prediction"`, or `"repeated_link_prediction"`. Required for the fit workflow.  In the load workflow, inferred automatically from the registered model if not provided.
- `eval_metric` (`str`, default: `None`) - Evaluation metric compatible with the chosen `task_type` (e.g. `"roc_auc"`, `"accuracy"`, `"rmse"`, `"link_prediction_precision@5"`).
- `use_current_time` (`bool`, default: `True`) - Use the current timestamp as the prediction time.  Default is `True`.
- `has_time_column` (`bool`, default: `None`) - Set to `True` when the task relationships use the `at` keyword for temporal ordering.  In the load workflow, inferred automatically from the registered model if not provided.
- `model_database` (`str`, default: `None`) - Snowflake database of a registered model (load workflow).
- `model_schema` (`str`, default: `None`) - Snowflake schema of a registered model (load workflow).
- `model_name` (`str`, default: `None`) - Name of the registered model (load workflow).
- `version_name` (`str`, default: `None`) - Version of the registered model (load workflow).
- `model_run_id` (`str`, default: `None`) - Run ID of a previously trained model (load workflow).
- `test_batch_size` (`int`, default: `None`) - Batch size used during inference.
- `stream_logs` (`bool`, default: `True`) - Stream training logs to stdout.  Default is `True`.
- `extract_embeddings` (`bool`, default: `False`) - Extract node embeddings during prediction.  Default is `False`.
- `dataset_alias` (`str`, default: `None`) - User chosen alias for the dataset.
- `parallel_reasoners_init` (`bool`, default: `True`) - Initialize the Predictive and Logic reasoners in parallel. Default is `True`.
- `**train_params` (`Any`, default: `{}`) - Additional hyperparameters forwarded to the GNN trainer (e.g. `n_epochs`, `lr`, `train_batch_size`, `device`, `head_layers`).  Ignored in the load workflow.

## Examples

Assuming the setup from the module-level Quick Start
([`relationalai.semantics.reasoners.predictive`](/api/python/v1.35/semantics/reasoners/predictive/index.md)):

```python
gnn = GNN(
    exp_database="EXPERIMENTS_DB",
    exp_schema="EXPERIMENTS_SCHEMA",
    graph=gnn_graph,
    property_transformer=property_transformer,
    source_concept=Students,
    train=Train,
    validation=Validation,
    task_type="binary_classification",
    eval_metric="roc_auc",
    device="cuda",
    n_epochs=5,
)
gnn.fit()
Students.predictions = gnn.predictions(domain=Test)
```

Load a registered model and predict:

```python
gnn = GNN(
    exp_database="EXPERIMENTS_DB",
    exp_schema="EXPERIMENTS_SCHEMA",
    model_database="MODELS_DB",
    model_schema="MODELS_SCHEMA",
    model_name="MY_MODEL",
    version_name="V1",
    source_concept=Students,
)
gnn.load()
Students.predictions = gnn.predictions(domain=Test)
```

## Methods

### .`fit`()

```python
GNN.fit() -> GNN
```

Train the GNN model.

Materializes graph and task tables, creates a trainer, and submits a
training job.  If training has already completed or is in progress,
calling `fit()` again is a no-op.

**Returns:**

- [`GNN`](/api/python/v1.35/semantics/reasoners/predictive/estimators/relational/gnn/gnn_class/index.md) - This instance, so that calls can be chained
  (e.g. `gnn.fit().predictions(domain=Test)`).

**Raises:**

- `ValueError` - If called in a load workflow.
- `NotImplementedError` - If deployments are enabled in the model configuration. Predictive
  reasoning is not yet supported in deployments.

### .`load`()

```python
GNN.load() -> GNN
```

Load a previously trained model for prediction.

Resolves the model from the Snowflake Model Registry (when
`model_database`, `model_schema`, `model_name`, and
`version_name` were provided) or by `model_run_id`, and
prepares the trainer for prediction.  If the model is already
loaded, calling `load()` again is a no-op.

**Returns:**

- [`GNN`](/api/python/v1.35/semantics/reasoners/predictive/estimators/relational/gnn/gnn_class/index.md) - This instance, so that calls can be chained
  (e.g. `gnn.load().predictions(domain=Test)`).

**Raises:**

- `ValueError` - If called in a fit workflow, or if the specified model or
  version does not exist.
- `NotImplementedError` - If deployments are enabled in the model configuration. Predictive
  reasoning is not yet supported in deployments.

### .`predictions`()

```python
GNN.predictions(domain: b.Relationship | b.Fragment | b.Chain) -> b.Relationship
```

Generate predictions on a test domain.

Materializes the test table, submits a prediction job, and returns a
[`Relationship`](/api/python/v1.35/semantics/frontend/base/relationship_class/index.md) that can be assigned to a
concept field for downstream querying.

The prediction attributes available on the returned relationship depend
on the task type:

- **Classification**: `.probs`, `.predicted_labels`
- **Regression**: `.predicted_value`
- **Link prediction**: `.rank`, `.scores`, `.predicted_<target>`

**Parameters:**

- `domain` (`Relationship or Fragment or Chain`) - The test split relationship (e.g. the `Test` relationship defined during data modeling).

**Returns:**

- `Relationship` - A prediction relationship to be assigned to the source concept
  (e.g. `User.predictions = gnn.predictions(domain=Test)`).

**Raises:**

- `TypeError` - If `domain` is not a Relationship, Fragment, or Chain.
- `ValueError` - If the model has not been fitted or loaded, or if the test
  domain schema does not match the training schema.
- `NotImplementedError` - If deployments are enabled in the model configuration. Predictive
  reasoning is not yet supported in deployments.

### .`score`()

```python
GNN.score(domain: Optional[b.Relationship | b.Fragment | b.Chain] = None) -> dict
```

Score the fitted or loaded model against a labeled domain.

Unlike [`GNN.predictions`](/api/python/v1.35/semantics/reasoners/predictive/estimators/relational/gnn/gnn_class/index.md#method-predictions), `score()` needs true labels to compare
against — pass a labeled split with the same shape as `train`
(source concept plus label column), or omit `domain` to score
against the `validation` split already configured on this instance.

Requires the same `graph` / `train` (node and edge tables) used to
fit the model, since scoring re-derives features the same way
[`GNN.fit`](/api/python/v1.35/semantics/reasoners/predictive/estimators/relational/gnn/gnn_class/index.md#method-fit) did.

**Parameters:**

- `domain` (`Relationship, Fragment, or Chain`, default: `None`) - A labeled split to score against. Defaults to the `validation` split passed to the constructor.

**Returns:**

- `dict` - Evaluation metrics computed by the backend for this task type.

**Raises:**

- `ValueError` - If no domain is available (neither passed nor configured via
  `validation=` at construction), if `graph`/`train` were not
  provided at construction, or if the model has not been fitted or
  loaded.
- `TypeError` - If `domain` is not a Relationship, Fragment, or Chain.
- `NotImplementedError` - If deployments are enabled in the model configuration.

### .`register_model`()

```python
GNN.register_model(
    model_database: str,
    model_schema: str,
    model_name: str,
    version_name: str,
    *,
    comment: Optional[str] = None
) -> None
```

Register a trained model in the Snowflake Model Registry.

After registration the model can be loaded in a later session by
passing the same `model_database`, `model_schema`,
`model_name`, and `version_name` to the [`GNN`](/api/python/v1.35/semantics/reasoners/predictive/estimators/relational/gnn/gnn_class/index.md) constructor.

**Parameters:**

- `model_database` (`str`) - Snowflake database for the model registry entry.
- `model_schema` (`str`) - Snowflake schema for the model registry entry.
- `model_name` (`str`) - Name under which to register the model.
- `version_name` (`str`) - Version label (e.g. `"v1"`).
- `comment` (`str`, default: `None`) - Free-text comment stored alongside the registry entry.

**Raises:**

- `ValueError` - If the model is already registered, or has not been fitted /
  loaded.

**Examples:**

```python
gnn.fit()
gnn.register_model(
    model_database="MODELS_DB",
    model_schema="MODELS_SCHEMA",
    model_name="STUDENT_CHURN",
    version_name="V1",
    comment="First baseline model",
)
```

### .`display_dataset_diagram`()

```python
GNN.display_dataset_diagram(show_dtypes: bool = False) -> None
```

Display the dataset schema inline in a Jupyter notebook.

Renders a diagram of the tables, columns, and foreign-key
relationships of the prepared dataset as an SVG in the current
notebook cell. Requires the Graphviz `dot` command-line tool
to be installed (`brew install graphviz` or
`apt install graphviz`).

**Parameters:**

- `show_dtypes` (`bool`, default: `False`) - Include column data types in the diagram. Default is `False`.

**Raises:**

- `ValueError` - If no dataset has been prepared yet (i.e. [`GNN.fit`](/api/python/v1.35/semantics/reasoners/predictive/estimators/relational/gnn/gnn_class/index.md#method-fit) has
  not been called).

**Examples:**

```python
gnn.fit()
gnn.display_dataset_diagram(show_dtypes=True)
```

### .`save_dataset_diagram`()

```python
GNN.save_dataset_diagram(path: str, show_dtypes: bool = False) -> None
```

Save the dataset schema diagram to a file.

Writes a diagram of the tables, columns, and foreign-key
relationships of the prepared dataset to *path*. The output
format is inferred from the file extension (e.g. `svg`,
`png`, `pdf`). Requires the Graphviz `dot` command-line
tool to be installed (`brew install graphviz` or
`apt install graphviz`).

**Parameters:**

- `path` (`str`) - Destination file path, e.g. `"schema.svg"` or `"schema.png"`.
- `show_dtypes` (`bool`, default: `False`) - Include column data types in the diagram. Default is `False`.

**Raises:**

- `ValueError` - If no dataset has been prepared yet (i.e. [`GNN.fit`](/api/python/v1.35/semantics/reasoners/predictive/estimators/relational/gnn/gnn_class/index.md#method-fit) has
  not been called).

**Examples:**

```python
gnn.fit()
gnn.save_dataset_diagram("schema.svg")
```

## Inheritance Hierarchy

```d2 inline layout=dagre theme=0 darkTheme=false pad=8 title="Inheritance hierarchy"
direction: down
"relationalai.semantics.reasoners.predictive.estimators.relational.gnn.GNN": {
  label: "GNN"
  link: "/api/python/v1.35/semantics/reasoners/predictive/estimators/relational/gnn/gnn_class"
}
"relationalai.semantics.reasoners.predictive.estimators.relational.base.PredictiveEstimator": {
  label: "PredictiveEstimator"
}
"abc.ABC": {
  label: "abc.ABC"
}
"relationalai.semantics.reasoners.predictive.estimators.relational.base.PredictiveEstimator" -> "relationalai.semantics.reasoners.predictive.estimators.relational.gnn.GNN"
"abc.ABC" -> "relationalai.semantics.reasoners.predictive.estimators.relational.base.PredictiveEstimator"
```

## Returned By

1. [semantics](/api/python/v1.35/semantics/index.md) / [reasoners](/api/python/v1.35/semantics/reasoners/index.md) / [predictive](/api/python/v1.35/semantics/reasoners/predictive/index.md) / [estimators](/api/python/v1.35/semantics/reasoners/predictive/estimators/index.md) / [relational](/api/python/v1.35/semantics/reasoners/predictive/estimators/relational/index.md) / [gnn](/api/python/v1.35/semantics/reasoners/predictive/estimators/relational/gnn/index.md)1) [GNN](/api/python/v1.35/semantics/reasoners/predictive/estimators/relational/gnn/gnn_class/index.md)1. [fit](/api/python/v1.35/semantics/reasoners/predictive/estimators/relational/gnn/gnn_class/index.md#method-fit)
      2. [load](/api/python/v1.35/semantics/reasoners/predictive/estimators/relational/gnn/gnn_class/index.md#method-load)