Skip to content
RelationalAI Docs

GNN

relationalai.semantics.reasoners.predictive
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 followed by GNN.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 followed by GNN.predictions.

In either workflow you can register the trained model in the Snowflake Model Registry using the GNN.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):

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:

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()

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 - 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()

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 - 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()

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 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()

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, 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 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()

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 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:

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()

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 has not been called).

Examples:

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

.save_dataset_diagram()

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 has not been called).

Examples:

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

Inheritance Hierarchy

GNNPredictiveEstimatorabc.ABC