GNN

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

GNN(
    *,
    exp_database: Optional[str] = None,
    exp_schema: Optional[str] = None,
    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,
    test: 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,
    task: Optional[b.Relationship] = None,
    name: 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.

In a deployed model the same GNN is model content: rai models deploy installs the tables it reads, and Students.predictions = gnn.predictions(domain=Test) declares a block that scores Test with a registered model on every refresh. gnn.fit() trains over the deployed tables and registers the result under model_name / version_name, the model the block scores. Every Snowflake location defaults to the deployment’s schema.

Parameters

• exp_database

  (str, default: None) - Snowflake database for experiment storage. Defaults to the deployment’s database in a deployed model.

• exp_schema

  (str, default: None) - Snowflake schema for experiment storage. Defaults to the deployment’s schema in a deployed model.

• 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

  (Relationship, default: None) - A labeled split held out from training and model selection, with the shape of train. In a deployed model the deployment installs a table for it, which score(domain=test) evaluates on.

• task

  (Relationship, default: None) - The labeled relationship the GNN predicts, declared without data: its first field is the source concept and its last the label, whose type the predicted labels take. A GNN that only scores has no training split to read these off; when train is given it must have the same shape.

• name

  (str, default: None) - Names this GNN in a deployed model. reasoners.predictive.gnns.<name> in raiconfig.yaml ties the name to a registered model and version per profile, which the GNN scores with in place of model_database, model_schema, model_name and version_name. A GNN that only scores must have one configured; a GNN that trains may have none.

• 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.

Notes:

In a deployed model, fit() exports nothing. It first deploys the model this GNN belongs to, which installs the train and validation tables, and returns at once when nothing installed changed (a new hyperparameter changes nothing installed). It then trains over those tables. When this GNN names a registered model, through reasoners.predictive.gnns or model_name / version_name, the run is registered there; otherwise GNN.register_model registers it.

Referenced By:

1. RelationalAI Documentation
   1. Build With RelationalAI
      1. Understand how PyRel works / Use advanced reasoning / Predictive reasoning
         1. Overview
            1. What predictive reasoning is
         2. Solve a classification problem
            1. Train the GNN and generate predictions
         3. Configure and Train a GNN
            1. Train the GNN
         4. Monitor training
            1. Find your experiment
         5. Register a model
            1. Register a model you just trained
         6. Use a pretrained model
         7. Understand GNN workflows
            1. Use the fit workflow

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

Notes:

In a deployed model, load() only checks that the registered model this GNN names exists. A deployed block loads that model itself on every refresh, so there is no trainer to prepare.

Referenced By:

1. RelationalAI Documentation
   1. Build With RelationalAI
      1. Understand how PyRel works / Use advanced reasoning / Predictive reasoning
         1. Register a model
            1. Register a previously trained model
         2. Use a pretrained model
            1. Load a model by run id
         3. Understand GNN workflows
            1. Use the load workflow

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

Notes:

In a deployed model this is a declaration rather than a job: it installs the tables the engine reads and a block that scores domain with the registered model on every refresh, and returns the same relationship, filled by that block. predict_df is the one-off that returns the rows as a DataFrame instead.

Referenced By:

1. RelationalAI Documentation
   1. Build With RelationalAI
      1. Understand how PyRel works / Use advanced reasoning / Predictive reasoning
         1. Overview
            1. What predictive reasoning is
         2. Solve a classification problem
            1. Train the GNN and generate predictions
         3. Make predictions
         4. Use a pretrained model
            1. Load a model by run id
         5. Understand GNN workflows
            1. Use the fit workflow
            2. Use the load workflow
            3. Score a fitted or loaded model

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

Notes:

In a deployed model, score() evaluates on what the deployment installed. Without domain, or with the validation split, it returns the best epoch’s validation metrics from this process’s fit() ({"val_metrics": ..., "best_epoch": ...}) and submits nothing. With the test split it submits a scoring job over the installed test table, for the run this process trained or else the registered model this GNN names, and returns the job’s result ({"test_metrics": ...}). The train split, and the validation split of a GNN that did not fit in this process, are scored the same way, so a GNN that GNN.load-ed a registered model scores any of the three.

Referenced By:

1. RelationalAI Documentation
   1. Build With RelationalAI
      1. Understand how PyRel works / Use advanced reasoning / Predictive reasoning
         1. Understand GNN workflows
            1. Score a fitted or loaded model

.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",
)

Referenced By:

1. RelationalAI Documentation
   1. Build With RelationalAI
      1. Understand how PyRel works / Use advanced reasoning / Predictive reasoning
         1. Register a model
            1. Register a model you just trained
            2. Register a previously trained model
         2. Use a pretrained model
            1. Load a model from the registry
         3. Understand GNN workflows
            1. Use the fit workflow

.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

Returned By

1. semantics / reasoners / predictive / estimators / relational / gnn
   1. GNN
      1. fit
      2. load

Referenced By

1. RelationalAI Documentation
   1. Build With RelationalAI
      1. Understand how PyRel works / Use advanced reasoning / Predictive reasoning
         1. Overview
            1. What predictive reasoning is
         2. Solve a classification problem
            1. Set up Snowflake access
            2. Import the libraries and create a model
            3. Train the GNN and generate predictions
         3. Construct a graph
            1. Pass the graph to the GNN
         4. Define a learning task
            1. Choose an evaluation metric
            2. Configure a temporal task
            3. Pass source and target concepts explicitly
         5. Encode features
         6. Configure and Train a GNN
            1. Example
            2. Train the GNN
         7. Monitor training
            1. Find your experiment
         8. Make predictions
         9. Register a model
            1. Register a model you just trained
            2. Register a previously trained model
         10. Use a pretrained model
             1. Configure a GNN in load mode
             2. Load a model by run id
             3. Load a model from the registry
         11. Understand GNN workflows
             1. Use the fit workflow
             2. Use the load workflow
             3. Score a fitted or loaded model