Embeddings API (preview)

Only available in Aura Graph Analytics

The gds.embedding endpoints are only available on Aura Graph Analytics sessions. They are not available on the GraphDataScience client used with the GDS plugin.
Preview feature

gds.embedding.create and gds.embedding.train are preview features. They may change or be renamed in future releases, and the client emits a UserWarning to remind you.

Node embeddings turn every node in a graph into a vector that captures where the node sits in the graph and what surrounds it. Nodes with similar neighborhoods get similar vectors, so you can feed the result straight into similarity search, clustering, classification, recommendations, or any model that expects vectors.

The gds.embedding endpoints are the shortest path from a graph to those vectors. You don’t need to pick an algorithm, tune hyperparameters, or know how graph machine learning works.

Endpoint What it does

gds.embedding.create

Generates embeddings for every node and adds them to your in-memory graph as a node property, named embedding unless you choose another name (mutate mode). With no encoder specified, it runs Fast Random Projection (FastRP) with sensible service-chosen defaults that are meant to work well across a wide range of graphs.

gds.embedding.train

Trains a GraphSAGE model on your own graph, using node properties and known classes, for when you want embeddings that take node properties into account or are tuned to one specific prediction task. The trained model is saved under a name you choose, like any other trained model in Aura Graph Analytics, and can be passed to gds.embedding.create. Store it with gds.model.store to reuse it from later sessions.

Requirements

  • GDS Python client 2.2 or later (pip install "graphdatascience>=2.2").

  • Aura API credentials (a client ID and secret), used to create and manage sessions.

  • An Aura Graph Analytics session. The gds.embedding endpoints are available on the client object returned by GdsSessions.get_or_create.

  • A graph in the session’s graph catalog, for example projected from your database or constructed from DataFrames.

Quick start

from graphdatascience.session import AuraAPICredentials, CloudLocation, GdsSessions

sessions = GdsSessions(api_credentials=AuraAPICredentials.from_env())  # reads CLIENT_ID, CLIENT_SECRET, and optionally PROJECT_ID
gds = sessions.get_or_create(
    session_name="embeddings",
    memory="2GB",
    cloud_location=CloudLocation(provider="<cloud-provider>", region="<cloud-region>"),
)
G = gds.graph.datasets.load_cora(undirected=True)  # or any graph in your session's catalog

result = gds.embedding.create(G) # adds an "embedding" property to every node of G
print(result.node_properties_written)  # 2708: one embedding per node

embeddings = gds.graph.node_properties.stream(G, "embedding")  # pandas DataFrame

This example starts a session that is not attached to an AuraDB instance, which needs a cloud_location. Replace the placeholders with a provider (gcp, aws, or azure) and a region available to your project; sessions.available_cloud_locations() lists them. For a session attached to an AuraDB instance, pass db_connection instead and leave cloud_location out.

The embedding property now exists on every node of the in-memory graph G. Stream it into pandas as shown above, or pass it to other GDS algorithms in the same session.

Using the embeddings

gds.embedding.create runs in mutate mode: it adds the embeddings to the in-memory graph G as a node property, named embedding unless you choose another name with mutate_property. Nothing is written to your database, and the graph isn’t copied. The property simply becomes part of G in your session’s graph catalog, so the embeddings are yours to use however you like:

# After gds.embedding.create(G) has run:

# Pull them into pandas
df = gds.graph.node_properties.stream(G, "embedding")

# Use them in other algorithms in the same session
similar = gds.knn.stream(G, node_properties="embedding", top_k=10)
clusters = gds.kmeans.stream(G, node_property="embedding", k=7)

The embeddings live as long as the graph does: when you drop the graph or the session ends, they are gone, so stream or write out whatever you want to keep.

To write them back to your Neo4j database, the graph must have been projected from that database, so that the session can map each node to its database counterpart:

gds.graph.node_properties.write(G, "embedding")
Calling create again with the same property name fails

If G already has a node property with the name you pass as mutate_property (embedding by default), create fails. To regenerate the embeddings under the same name, for example after changing a setting or as part of a scheduled pipeline, drop the property first:

gds.graph.node_properties.drop(G, ["embedding"], fail_if_missing=False)
gds.embedding.create(G)

To keep both sets instead, give each call its own mutate_property.

Create embeddings

gds.embedding.create(G) is the whole call. With no graph_encoder, create uses FastRP with settings chosen by the service. With no mutate_property, it stores the embeddings as the embedding node property of G. To change either, pass them explicitly, as in the examples below.

Parameters

Name Type Default Description

G

Graph

required

The graph to embed.

graph_encoder

FastRPConfig, str, or None

None

What produces the embeddings: a FastRPConfig, or the name of a model saved by gds.embedding.train. Leave it out (None) to use FastRP with service-chosen settings, which is the recommended starting point.

embedding_dimension

int

None

Dimension of the embedding vectors. If unset, the service picks one.

mutate_property

str

"embedding"

Name of the node property that holds the embeddings on the in-memory graph G. If G already has a property with this name, the call fails.

random_seed

int

None

Seed for reproducible results. If unset, a seed is selected at random.

node_labels

list[str]

["*"]

Only nodes with one of these labels are included.

relationship_types

list[str]

["*"]

Only relationships of these types are included.

input_properties

list[str]

[]

Node properties to use as input features. Required when graph_encoder is a trained model (pass the properties it was trained on). Ignored for FastRP, which only uses the graph structure.

job_id

str

None

Identifier for the computation.

Result

gds.embedding.create returns an EmbeddingCreateResult.

Field Type Description

node_properties_written

int

Number of nodes that received an embedding.

compute_millis

int

Time spent computing the embeddings, in milliseconds.

pre_processing_millis

int

Time spent preparing the computation, in milliseconds.

mutate_millis

int

Time spent adding the property to the in-memory graph, in milliseconds.

configuration

dict

The configuration the job ran with.

Examples

Choose the embedding dimension and make the run reproducible. G already has an embedding property from the quick start, and create fails if the property exists, so drop it first:

gds.graph.node_properties.drop(G, ["embedding"], fail_if_missing=False)
gds.embedding.create(
    G,
    embedding_dimension=256,
    random_seed=42,
)

If you leave embedding_dimension unset, the service picks a dimension for you. The same applies to the other FastRP settings (coefficients and normalization_strength), which are set through FastRPConfig.

Keep the earlier embeddings by storing the new ones under a different property name:

gds.embedding.create(G, mutate_property="embedding_256", embedding_dimension=256)

Restrict the graph to the node labels and relationship types you want, and keep the result in a property of its own so a full-graph embedding property stays as it is:

gds.embedding.create(
    G,
    mutate_property="embedding_subset",
    node_labels=["Paper"],
    relationship_types=["CITES"],
)

On a graph with several node labels or relationship types, list only the ones you want to include; the embeddings still cover every node of the (filtered) graph, so you may get vectors for more labels than you plan to use. Keep a label when it is needed for connectivity or for context: if nodes of one label are only connected to the rest of the graph through nodes of another, dropping the latter would disconnect them.

To get the embeddings out of the session, see Using the embeddings.

Train your own model

gds.embedding.create is designed to give good results without training, using only the structure of the graph. Use gds.embedding.train when you want embeddings that take node properties into account, which is currently the only way to do so, or that are tuned to a specific task such as node classification.

Training fits a GraphSAGE encoder so that its embeddings are good at predicting a node property that holds your known classes, and saves it under the name you give it. You then pass that name as the graph_encoder of gds.embedding.create.

G = gds.graph.datasets.load_cora(undirected=True)

gds.embedding.train(
    G,
    model_name="paper-topic-encoder",
    target_label="Paper",
    target_property="subject",
    input_properties=["features"],
)

gds.embedding.create(
    G,
    graph_encoder="paper-topic-encoder",
    input_properties=["features"],
)

To tune the encoder, pass a GraphSAGEConfig as graph_encoder.

Parameters

Name Type Default Description

G

Graph

required

The graph to train on.

model_name

str

required

Name to save the trained model under.

target_label

str

required

Node label to train on.

target_property

str

required

Node property holding the known classes to train on.

input_properties

list[str]

required

Node properties the encoder learns from.

graph_encoder

GraphSAGEConfig or None

None

Configuration of the GraphSAGE encoder to train. If unset, the service trains a GraphSAGE encoder with default settings.

embedding_dimension

int

None

Dimension of the embedding vectors. If unset, the service picks one.

decoder

MLPClassifierConfig or None

None

Configuration of the classifier trained on top of the encoder’s embeddings. Most users can leave it out.

num_epochs

int

None

Maximum number of training epochs. If unset, the service chooses.

batch_size

int

None

Number of examples per training batch. If unset, the service chooses.

random_seed

int

None

Seed for reproducible results. If unset, a seed is selected at random.

node_labels

list[str]

["*"]

Only nodes with one of these labels are included.

relationship_types

list[str]

["*"]

Only relationships of these types are included.

job_id

str

None

Identifier for the computation.

Result

gds.embedding.train returns an EmbeddingTrainResult with compute_millis, pre_processing_millis, and configuration.

Store and reuse trained models

A model trained with gds.embedding.train is a regular Aura Graph Analytics model. It is listed in the model catalog next to models from other training procedures, and you manage it with the same gds.model endpoints.

Storing a model is optional. To keep it for later sessions, persist it with gds.model.store. Stored models live in the Aura Graph Analytics model catalog, a persistent store scoped to your Aura project, the cloud provider, and region. They remain available after the session that trained them has ended.

The other model catalog endpoints work the same way as for any model.

Encoder and decoder configurations

All configuration classes are imported from graphdatascience.embedding:

from graphdatascience.embedding import (
    FastRPConfig,
    GraphSAGEConfig,
    MLPClassifierConfig,
)
Configuration Fields

FastRPConfig

out_dim, coefficients, normalization_strength. All optional.

GraphSAGEConfig

hidden_dims, num_neighbors, out_dim, dropout. All optional. num_neighbors must have exactly one more entry than hidden_dims.

MLPClassifierConfig

hidden_dims, dropout. A multi-layer perceptron classifier. If hidden_dims is empty ([]), the classifier has no hidden layers and acts as a linear classifier.

Optional fields you leave out are left to the service.

Good to know

  • Embeddings cover the whole graph. Each call computes embeddings for all nodes in the (filtered) graph. To embed new nodes, add them to the graph and call create again.

  • Seeds make runs reproducible. Set random_seed to get the same embeddings from the same graph across runs. If the graph changes, the embeddings can change too, even with the same seed.

  • Trained models have seen your known classes. Metrics computed on the nodes a model was trained on are optimistic by construction.

  • Store models you want to keep. To use a trained model from a later session, persist it with gds.model.store. Stored models are scoped to your project and cloud location. Delete the ones you no longer need with gds.model.delete.

  • Sessions bill while they run. Delete the session when you’re done using gds.delete().

Next steps