Embedding API (preview)

This Jupyter notebook is hosted here in the Neo4j Graph Data Science Client Github repository.

The notebook shows how to use the graphdatascience Python library to compute node embeddings with the Embedding API, a unified interface for creating and training embedding models in a GDS Session.

The Embedding API is a preview feature and may change or be removed in future releases.

We consider the Cora citation network, loaded with the built-in dataset tooling. We first compute FastRP embeddings without any training, and then train a GraphSAGE encoder with an MLP classifier decoder and apply the trained model to generate embeddings.

Prerequisites

This notebook requires having the Aura Graph Analytics feature enabled for your Neo4j Aura project, as well as Aura API credentials. The credentials are read from environment variables, optionally loaded from a sessions.env file (see sessions.env.template).

You also need to have the graphdatascience Python library installed, version 2.0a7 or later.

%pip install "graphdatascience>=2.0a7" python-dotenv

import os

from dotenv import load_dotenv

load_dotenv("sessions.env")

Creating a GDS Session

The entry point for managing GDS Sessions is the GdsSessions object, which requires Aura API credentials. Since this notebook works purely on in-memory graphs, the session does not need to be connected to a database.

from graphdatascience.session import AuraAPICredentials, CloudLocation, GdsSessions, SessionMemory

# Create a new GdsSessions object
sessions = GdsSessions(
    api_credentials=AuraAPICredentials(
        os.environ.get("CLIENT_ID"),
        os.environ.get("CLIENT_SECRET"),
        os.environ.get("PROJECT_ID"),
    )
)

gds = sessions.get_or_create(
    session_name="my_session",
    memory=SessionMemory.m_2GB,
    cloud_location=CloudLocation(provider="gcp", region="europe-west1"),
)

Embeddings

The simplest way to compute node embeddings is gds.embedding.create with a default graph encoder, here FastRP, which requires no training. The resulting embeddings are added to the in-memory graph which we then stream back as a DataFrame.

with gds.graph.datasets.load_cora(undirected=True) as G:
    create_result = gds.embedding.create(G)
    node_embeddings = gds.graph.node_properties.stream(G, node_properties="embedding")
node_embeddings.head()

Training a GraphSAGE model

gds.embedding.train trains a graph encoder together on a specific task (classifying a target property). Here we have been specific about the graph encoder, its hyperparameters and training parameters. Left out optional parameters will be replaced by defaults consistent with the specified parameters.

The trained model is saved under a specified name, and can then be used to generate embeddings using gds.embedding.create.

from graphdatascience.embedding import GraphSAGEConfig

with gds.graph.datasets.load_cora(undirected=True) as G:
    train_result = gds.embedding.train(
        G=G,
        graph_encoder=GraphSAGEConfig(hidden_dims=[128], num_neighbors=[10, 5], out_dim=64),
        model_name="cora_model",
        target_label="Paper",
        target_property="subject",
        num_epochs=25,
        batch_size=64,
        input_properties=["features"],
    )
    create_result = gds.embedding.create(
        G=G,
        graph_encoder="cora_model",
        input_properties=["features"],
        mutate_property="my_graphsage_embedding",
    )

    graphsage_embeddings = gds.graph.node_properties.stream(G, node_properties="my_graphsage_embedding")
graphsage_embeddings.head()

Cleanup

After the analysis is done, we delete the trained model and the session. As this example is not connected to a Neo4j DB, make sure to persist any results you want to keep beforehand. Deleting the session releases all resources associated with it, and stops incurring costs.

gds.model.delete("cora_model")
gds.model.delete("hyperparameter_tuned_cora_model")

sessions.delete(session_name="my_session")