Skip to content

Pretrained Foundation Models

K3 Node provides native implementations of modern Graph Foundation Models with built-in checkpoint downloaders and weight loaders for official pre-trained models.

This page covers general-purpose graph foundation models. For domain-specific pretrained checkpoints see the Materials (CHGNet, M3GNet, MEGNet, TensorNet, ...) and Bio (Uni-Mol docking) API pages. See the Fine-Tuning Recipes guide for end-to-end worked examples of adapting any of these to a downstream task.

1. GraphMAE2 (Masked Autoencoder for Graphs)

GraphMAE2 is an advanced masked autoencoder for self-supervised graph representation learning featuring multi-task reconstruction (scaled cosine error) and projection heads.

Architecture

k3_node.models.graphmae2.GraphMAE2

Bases: Layer

The GraphMAE2 model from "GraphMAE2: A Decoding-Enhanced Masked Self-Supervised Learning Framework for Graphs" <https://arxiv.org/abs/2304.04779>_.

Parameters:

Name Type Description Default
in_dim int

Dimensionality of input node features.

 required
num_hidden int

Dimensionality of hidden node representations.

 required
num_layers int

Number of encoder layers. (default: 4)

 4
num_dec_layers int

Number of decoder layers. (default: 1)

 1
num_remasking int

Number of remasking views in decoder. (default: 3)

 3
nhead int

Number of attention heads in encoder. (default: 8)

 8
nhead_out int

Number of attention heads in decoder output. (default: 1)

 1
activation str

Activation function. (default: "prelu")

 'prelu'
feat_drop float

Node feature dropout rate. (default: 0.2)

 0.2
attn_drop float

Attention dropout rate. (default: 0.1)

 0.1
negative_slope float

LeakyReLU negative slope. (default: 0.2)

 0.2
residual bool

Whether to use residual connections. (default: True)

 True
norm str

Normalization type ("layernorm", "batchnorm", or None). (default: "layernorm")

 'layernorm'
mask_rate float

Fraction of input nodes to mask. (default: 0.5)

 0.5
remask_rate float

Fraction of latent nodes to remask. (default: 0.5)

 0.5
remask_method str

Remasking method ("random" or "fixed"). (default: "random")

 'random'
loss_fn str

Reconstruction loss type ("sce" or "mse"). (default: "sce")

 'sce'
alpha_l float

Power exponent in Scaled Cosine Error loss. (default: 2.0)

 2.0
lam float

Weight of the latent prediction loss term. (default: 1.0)

 1.0
momentum float

Teacher EMA update momentum. (default: 0.996)

 0.996
delayed_ema_epoch int

Epoch to begin EMA teacher updates. (default: 0)

 0
Example 
import numpy as np
from k3_node.models import GraphMAE2

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

model = GraphMAE2(in_dim=8, num_hidden=32, num_layers=2, num_dec_layers=1, nhead=4, nhead_out=1)
print(tuple(model(x, edge_index).shape))  # (10, 32): node embeddings
loss = model.loss(x, edge_index)  # masked feature reconstruction loss for pre-training
print(tuple(loss.shape))  # ()

call(x, edge_index, training=False)

Forward pass: returns node embeddings by default.

ema_update(momentum=None)

Updates teacher EMA parameters.

embed(x, edge_index, training=None)

Generates node embeddings with the encoder.

encoding_mask_noise(x, mask_rate=None, mask_nodes=None)

Masks node features for encoder input.

from_pretrained(dataset='ogbn-arxiv', folder='checkpoints', download=True, **kwargs) classmethod

Instantiates a GraphMAE2 model with pre-trained weights downloaded from Google Drive.

Parameters:

Name Type Description Default
dataset str

Dataset name ("ogbn-arxiv", "ogbn-products", "mag-scholar-f", or "ogbn-papers100M").

 'ogbn-arxiv'
folder str

Directory to store/find checkpoints. (default: "checkpoints")

 'checkpoints'
download bool

Whether to download checkpoint if missing locally. (default: True)

 True
**kwargs

Overrides for model hyperparameters.

 {}

Returns:

Name Type Description
GraphMAE2 GraphMAE2

Model instance loaded with pre-trained weights.

load_weights_from_checkpoint(checkpoint_path=None, dataset=None, folder='checkpoints', download=True)

Loads weights from a PyTorch state dict checkpoint or Google Drive.

loss(x, edge_index, mask_nodes=None, targets=None, epoch=0, training=True)

Computes GraphMAE2 loss: attribute reconstruction loss + latent prediction loss.

random_remask(rep, remask_rate=None, remask_nodes=None)

Remasks latent representation for decoder input.

Checkpoint Loading

k3_node.models.graphmae2.load_graphmae2_weights(model, checkpoint_path=None, dataset=None, folder='checkpoints', download=True)

Loads pre-trained weights from a GraphMAE2 PyTorch checkpoint (.pt).

If the checkpoint does not exist locally and download=True, it will be automatically downloaded from the official Google Drive folder using download_google_url.

Parameters:

Name Type Description Default
model GraphMAE2

The target GraphMAE2 model instance.

 required
checkpoint_path str

Local path to .pt file or dataset name.

 None
dataset str

Dataset name if downloading from Google Drive.

 None
folder str

Directory to store downloaded checkpoints. (default: "checkpoints")

 'checkpoints'
download bool

Whether to download checkpoint if missing locally. (default: True)

 True

Checkpoint Downloader

k3_node.models.graphmae2.download_graphmae2_checkpoint(dataset, folder='checkpoints', log=True)

Downloads a pre-trained GraphMAE2 checkpoint from Google Drive using download_google_url.

Google Drive folder: https://drive.google.com/drive/folders/1GiuP0PtIZaYlJWIrjvu73ZQCJGr6kGkh

Parameters:

Name Type Description Default
dataset str

Name of dataset ("ogbn-arxiv", "ogbn-products", "mag-scholar-f", or "ogbn-papers100M").

 required
folder str

Target directory to save the checkpoint. (default: "checkpoints")

 'checkpoints'
log bool

Whether to print download progress. (default: True)

 True

Returns:

Name Type Description
str str

Absolute path to the downloaded checkpoint file.

Usage Example

from k3_node.models.graphmae2 import GraphMAE2, load_graphmae2_weights, download_graphmae2_checkpoint

# Download checkpoint (one of "ogbn-arxiv", "ogbn-products", "mag-scholar-f", "ogbn-papers100M")
ckpt_path = download_graphmae2_checkpoint("ogbn-arxiv")

# Initialize and load model
model = GraphMAE2(in_dim=1433, num_hidden=512, out_dim=1433, num_layers=2)
load_graphmae2_weights(model, ckpt_path)

2. Graphormer & Graphormer3D

Graphormer is a Graph Transformer with spatial, degree, and edge attention biases. Graphormer3D extends this architecture with 3D Gaussian RBF distance biases for molecular conformation and quantum chemistry.

Graphormer (2D)

k3_node.models.graphormer.Graphormer

Bases: Model

Graphormer model for molecular graph representation and property prediction from "Do Transformers Really Perform Badly for Graph Representation?" <https://arxiv.org/abs/2106.05234>_.

Parameters:

Name Type Description Default
num_atoms int

Maximum atom vocabulary size. (default: 512)

 512
num_in_degree int

Maximum in-degree. (default: 512)

 512
num_out_degree int

Maximum out-degree. (default: 512)

 512
num_edges int

Maximum edge vocabulary size. (default: 512)

 512
num_spatial int

Maximum spatial distance. (default: 512)

 512
num_edge_dis int

Maximum edge distance multiplier. (default: 128)

 128
edge_type str

Type of edge encoding ("multi_hop" or "single_hop"). (default: "multi_hop")

 'multi_hop'
multi_hop_max_dist int

Max distance for multi-hop paths. (default: 20)

 20
num_encoder_layers int

Number of Transformer layers. (default: 12)

 12
embedding_dim int

Hidden embedding dimension. (default: 768)

 768
ffn_embedding_dim int

FFN intermediate dimension. (default: 768)

 768
num_attention_heads int

Number of attention heads. (default: 32)

 32
dropout float

Dropout rate. (default: 0.0)

 0.0
attention_dropout float

Attention dropout rate. (default: 0.1)

 0.1
activation_dropout float

Activation dropout rate. (default: 0.1)

 0.1
encoder_normalize_before bool

Whether to apply LayerNorm before encoder blocks. (default: True)

 True
pre_layernorm bool

Whether to use Pre-LN. (default: False)

 False
num_classes int

Output dimension for prediction head. (default: 1)

 1
activation_fn str

Activation function. (default: "gelu")

 'gelu'
**kwargs

Additional model arguments.

 {}
Example 
import numpy as np
from k3_node.models import Graphormer

# A batch of 2 graphs padded to 4 nodes, preprocessed into Graphormer's dense inputs
data = {
    "x": np.random.randint(1, 15, size=(2, 4, 2)),  # 2 categorical features per node
    "in_degree": np.random.randint(0, 7, size=(2, 4)),
    "out_degree": np.random.randint(0, 7, size=(2, 4)),
    "attn_bias": np.zeros((2, 5, 5), dtype="float32"),  # +1 for the virtual graph token
    "spatial_pos": np.random.randint(0, 7, size=(2, 4, 4)),  # shortest-path distances
    "edge_input": np.random.randint(0, 7, size=(2, 4, 4, 2, 2)),  # edge types along shortest paths
}

model = Graphormer(num_atoms=16, num_in_degree=8, num_out_degree=8, num_edges=8, num_spatial=8,
                   num_edge_dis=4, num_encoder_layers=2, embedding_dim=32, ffn_embedding_dim=64,
                   num_attention_heads=4, num_classes=1)
out = model(data)  # one prediction per graph
print(tuple(out.shape))  # (2, 1)

call(batched_data=None, x=None, in_degree=None, out_degree=None, attn_bias=None, spatial_pos=None, edge_input=None, attn_edge_type=None, perturb=None, return_all=False, training=False)

Forward pass for Graphormer.

Accepts either a single dictionary batched_data containing graph tensors, or individual tensor arguments.

Returns:

Type Description

Tensor or Tuple[Tensor, Tensor]: Graph prediction tensor of shape [batch_size, num_classes],

or if return_all=True, a tuple of (graph_pred, all_node_features).

embed(batched_data=None, x=None, in_degree=None, out_degree=None, attn_bias=None, spatial_pos=None, edge_input=None, attn_edge_type=None)

Computes graph and node embeddings without applying the prediction head.

from_pretrained(pretrained_name='pcqm4mv1_graphormer_base', folder='checkpoints', download=True, **kwargs) classmethod

Instantiates a Graphormer model with pre-trained weights.

Graphormer3D (3D)

k3_node.models.graphormer_3d.Graphormer3D

Bases: Model

Graphormer-3D model for 3D molecular structure modeling, energy, and force prediction from "Benchmarking Graphormer on Large-Scale Molecular Modeling Datasets" <https://arxiv.org/abs/2203.04810>_.

Parameters:

Name Type Description Default
layers int

Number of encoder layers per block. (default: 12)

 12
blocks int

Number of repeated encoder blocks. (default: 4)

 4
embed_dim int

Hidden embedding dimension. (default: 768)

 768
ffn_embed_dim int

FFN intermediate dimension. (default: 768)

 768
attention_heads int

Number of attention heads. (default: 48)

 48
num_kernel int

Number of Gaussian basis kernels. (default: 128)

 128
atom_types int

Number of atom types. (default: 64)

 64
dropout float

Dropout probability. (default: 0.1)

 0.1
attention_dropout float

Attention dropout. (default: 0.1)

 0.1
activation_dropout float

FFN activation dropout. (default: 0.0)

 0.0
input_dropout float

Input features dropout. (default: 0.0)

 0.0
**kwargs

Additional model arguments.

 {}
Example 
import numpy as np
from k3_node.models import Graphormer3D

atoms = np.array([[1, 2, 3, 4, 0], [2, 3, 4, 0, 0]])  # atom types, 0 = padding
tags = np.array([[1, 1, 2, 2, 0], [1, 2, 2, 0, 0]])  # e.g. surface/adsorbate tags (OC20)
pos = np.random.rand(2, 5, 3).astype("float32")
model = Graphormer3D(layers=2, blocks=2, embed_dim=32, ffn_embed_dim=64, attention_heads=4,
                     num_kernel=16, atom_types=16)
energy, forces = model(atoms, tags, pos)
print(tuple(energy.shape), tuple(forces.shape))  # (2,) (2, 5, 3)

call(atoms, tags, pos, real_mask=None, training=False)

Forward pass for Graphormer-3D predicting total energy and atomic forces.

Parameters:

Name Type Description Default
atoms Tensor

Atom indices of shape [batch_size, num_nodes].

 required
tags Tensor

Tag indices of shape [batch_size, num_nodes] (0: fixed, 1: sub-surface, 2: surface).

 required
pos Tensor

3D atomic coordinates of shape [batch_size, num_nodes, 3].

 required
real_mask Tensor

Valid non-padding mask of shape [batch_size, num_nodes]. If None, non-zero atom indices are considered valid.

 None
training bool

Training flag. (default: False)

 False

Returns:

Type Description

Tuple[Tensor, Tensor]: Tuple of predicted energy [batch_size] and atomic forces

[batch_size, num_nodes, 3].

from_pretrained(pretrained_name='oc20is2re_graphormer3d_base', folder='checkpoints', download=True, **kwargs) classmethod

Instantiates a Graphormer3D model with pre-trained weights.

Checkpoint Utilities

k3_node.models.graphormer.load_graphormer_weights(model, checkpoint_path=None, pretrained_name=None, folder='checkpoints', download=True)

Loads weights from a PyTorch (.pt or .bin) checkpoint into a Keras Graphormer model.

Parameters:

Name Type Description Default
model Graphormer

Target model instance.

 required
checkpoint_path str

Path to .pt or .bin file.

 None
pretrained_name str

Name of pre-trained model to load or download.

 None
folder str

Checkpoints directory. (default: "checkpoints")

 'checkpoints'
download bool

Whether to download if missing. (default: True)

 True

Returns:

Name Type Description
Graphormer Graphormer

The model with loaded weights.

k3_node.models.graphormer.download_graphormer_checkpoint(name, folder='checkpoints', log=True)

Downloads a pre-trained Graphormer checkpoint to the specified folder.

Parameters:

Name Type Description Default
name str

Pretrained checkpoint name (e.g., "pcqm4mv1_graphormer_base", "pcqm4mv2_graphormer_base").

 required
folder str

Destination folder. (default: "checkpoints")

 'checkpoints'
log bool

Whether to print download progress. (default: True)

 True

Returns:

Name Type Description
str str

Absolute path to the downloaded file.

3. GraphGPS (General Powerful Scalable Graph Transformer)

GraphGPS is a hybrid graph transformer architecture combining local message-passing neural networks (CustomGatedGCN) and global linear/full attention with Random Walk Structural Encodings (RWSE).

GPSModel

k3_node.models.gps_model.GPSModel

Bases: Model

GraphGPS: General Powerful Scalable Graph Transformer from the "Recipe for a General, Powerful, Scalable Graph Transformer" <https://arxiv.org/abs/2205.12454>_ paper (NeurIPS 2022).

Parameters:

Name Type Description Default
dim_in int

Initial input feature dimension. (default: 256)

 256
dim_out int

Target output dimension. (default: 1)

 1
num_layers int

Number of GPS layers. (default: 16)

 16
dim_hidden int

Hidden embedding dimension. (default: 256)

 256
num_heads int

Number of attention heads. (default: 8)

 8
local_gnn_type str

Local MPNN layer type. (default: "CustomGatedGCN")

 'CustomGatedGCN'
act str

Activation function. (default: "gelu")

 'gelu'
dropout float

Dropout probability. (default: 0.1)

 0.1
attn_dropout float

Attention dropout probability. (default: 0.1)

 0.1
batch_norm bool

Whether to use batch normalization. (default: True)

 True
layer_norm bool

Whether to use layer normalization. (default: False)

 False
node_encoder_type str

Node encoder type ("Atom+RWSE", "Atom", "Linear", or None). (default: "Atom+RWSE")

 'Atom+RWSE'
edge_encoder_type str

Edge encoder type ("Bond", "Linear", or None). (default: "Bond")

 'Bond'
atom_feature_dims List[int]

Categorical feature vocabulary sizes for atom features.

 None
bond_feature_dims List[int]

Categorical feature vocabulary sizes for bond features.

 None
rwse_num_steps int

Number of RWSE steps. (default: 16)

 16
rwse_dim_pe int

RWSE embedding dimension. (default: 20)

 20
graph_pooling str

Graph pooling type ('mean', 'add', 'max'). (default: "mean")

 'mean'
head_layers int

Number of hidden layers in prediction head. (default: 2)

 2
**kwargs

Additional model arguments.

 {}
Example 
import numpy as np
from k3_node.models import GPSModel

# OGB-style integer atom and bond features (9 per atom, 3 per bond)
x = np.random.randint(0, 2, size=(5, 9))
edge_index = np.array([[0, 1, 1, 2, 3, 4], [1, 0, 2, 1, 4, 3]])
edge_attr = np.random.randint(0, 2, size=(6, 3))
batch = np.array([0, 0, 0, 1, 1])  # two graphs
rwse = np.random.rand(5, 16).astype("float32")  # random-walk structural encodings

model = GPSModel(dim_in=32, dim_out=1, num_layers=2, dim_hidden=32, num_heads=4,
                 node_encoder_type="Atom+RWSE", edge_encoder_type="Bond",
                 rwse_num_steps=16, rwse_dim_pe=8)
pred = model(x, edge_index, edge_attr=edge_attr, pestat_RWSE=rwse, batch=batch, batch_size=2)
print(tuple(pred.shape))  # (2, 1): one prediction per graph

Checkpoint Utilities

k3_node.models.gps_model.load_gps_weights(model, checkpoint_path)

Loads trained PyTorch GraphGPS checkpoint weights into a Keras 3 GPSModel.

Parameters:

Name Type Description Default
model GPSModel

Target GPSModel instance.

 required
checkpoint_path str

Path to PyTorch .ckpt or .pt checkpoint file.

 required

Returns:

Name Type Description
GPSModel

The model with loaded weights.

k3_node.models.gps_model.download_gps_checkpoint(checkpoint_name='pcqm4m-GPS+RWSE.deep', cache_dir=None)

Downloads and extracts a pretrained GraphGPS checkpoint.

Parameters:

Name Type Description Default
checkpoint_name str

Name of the checkpoint. Currently supported: "pcqm4m-GPS+RWSE.deep".

 'pcqm4m-GPS+RWSE.deep'
cache_dir str

Cache directory to store downloaded checkpoint.

 None

Returns:

Name Type Description
str str

Absolute path to the extracted .ckpt checkpoint file.

Usage Example

from k3_node.models.gps_model import GPSModel, load_gps_weights, download_gps_checkpoint

# Download official pcqm4m-GPS+RWSE.deep checkpoint
ckpt_path = download_gps_checkpoint("pcqm4m-GPS+RWSE.deep")

# Initialize and load weights
model = GPSModel(dim_in=256, dim_out=1, num_layers=16, dim_hidden=256, num_heads=8)
load_gps_weights(model, ckpt_path)

4. GROVER (Graph Representation from Self-Supervised Message Passing Transformer)

GROVER incorporates dual-track message passing across directed bonds and atoms with transformer self-attention and scope-based molecular readouts.

GROVER

k3_node.models.grover.GROVER

Bases: Model

Complete GROVER Model.

Example 
import numpy as np
from k3_node.models import GROVER

# Chemprop-style molecule batch; index 0 is a zero padding atom/bond.
f_atoms = np.random.rand(5, 151).astype("float32")  # atom features
f_bonds = np.random.rand(7, 165).astype("float32")  # (directed) bond features
f_atoms[0], f_bonds[0] = 0.0, 0.0
a2b = np.array([[0, 0], [2, 0], [1, 4], [3, 0], [0, 0]])  # incoming bonds of each atom
b2a = np.array([0, 1, 2, 2, 3, 1, 3])  # source atom of each bond
b2revb = np.array([0, 2, 1, 4, 3, 6, 5])  # reverse bond of each bond
a_scope = np.array([[1, 2], [3, 2]])  # (start, size) of each molecule's atoms
b_scope = np.array([[1, 3], [4, 3]])
a2a = b2a[a2b]  # neighboring atoms
batch = (f_atoms, f_bonds, a2b, b2a, b2revb, a_scope, b_scope, a2a)

model = GROVER(hidden_size=32, edge_fdim=165, node_fdim=151, num_mt_block=1, num_attn_head=2,
               depth=3, atom_emb_output="both", readout_type="mean")
out = model(batch)
print(tuple(out["atom_from_atom"].shape))  # (5, 32): atom embeddings
print(tuple(model.get_fingerprint(batch, fingerprint_source="both").shape))  # (2, 128): one per molecule

call(inputs, training=False)

Inputs can be a tuple/list: (f_atoms, f_bonds, a2b, b2a, b2revb, a_scope, b_scope, a2a).

get_fingerprint(inputs, fingerprint_source='both', features_batch=None, training=False)

Generate molecule-level fingerprints using Readout.

Checkpoint Utilities

k3_node.models.grover.load_grover_weights(model, checkpoint_path)

Loads PyTorch GROVER checkpoint state dict into Keras 3 GROVER model.

k3_node.models.grover.download_grover_checkpoint(checkpoint_name='grover_base', cache_dir=None)

Download official GROVER pre-trained model checkpoint from Google Drive.

Usage Example

from k3_node.models.grover import GROVER, load_grover_weights, download_grover_checkpoint

ckpt_path = download_grover_checkpoint("grover_base")

model = GROVER(hidden_size=800, num_attn_head=4, depth=6, num_mt_block=1)
load_grover_weights(model, ckpt_path)

5. Mole-BERT (Self-Supervised Molecular GNN)

Mole-BERT pre-trains a 5-layer GIN backbone with categorical atom and bond embeddings, batch normalization, and Jumping Knowledge for downstream molecular property prediction.

MoleBERT

k3_node.models.mole_bert.MoleBERT

Bases: Model

Complete Mole-BERT Model with graph-level pooling and property prediction head.

Example 
import numpy as np
from k3_node.models import MoleBERT

x = np.stack([np.random.randint(0, 119, size=5), np.random.randint(0, 3, size=5)], axis=1)  # atom type, chirality
edge_index = np.array([[0, 1, 1, 2, 3, 4], [1, 0, 2, 1, 4, 3]])
edge_attr = np.stack([np.random.randint(0, 5, size=6), np.random.randint(0, 3, size=6)], axis=1)  # bond type, direction
batch = np.array([0, 0, 0, 1, 1])  # two molecules

model = MoleBERT(num_layer=3, emb_dim=32, num_tasks=2, graph_pooling="mean")
logits, node_rep = model((x, edge_index, edge_attr, batch))
print(tuple(logits.shape), tuple(node_rep.shape))  # (2, 2) (5, 32): per-molecule predictions, per-atom embeddings

call(inputs, training=False)

Call MoleBERT model.

inputs can be a tuple: (x, edge_index, edge_attr) or (x, edge_index, edge_attr, batch).

Checkpoint Utilities

k3_node.models.mole_bert.load_mole_bert_weights(model, checkpoint_path)

Loads official PyTorch Mole-BERT.pth checkpoint state dict into Keras 3 MoleBERT model.

k3_node.models.mole_bert.download_mole_bert_checkpoint(cache_dir=None)

Downloads official Mole-BERT.pth checkpoint from GitHub.

Usage Example

from k3_node.models.mole_bert import MoleBERT, load_mole_bert_weights, download_mole_bert_checkpoint

ckpt_path = download_mole_bert_checkpoint()

model = MoleBERT(num_layer=5, emb_dim=300, num_tasks=1)
load_mole_bert_weights(model, ckpt_path)

6. Uni-Mol Family (3D Molecular Pretraining)

Uni-Mol pretrains a 3D-coordinate-aware transformer on large-scale molecular conformations for downstream property prediction, conformation generation, and (see Bio API) protein–ligand docking. UniMol2 and UniMol-Plus are scaled-up successors with two-track atom/pair transformers.

UniMolModel

k3_node.models.unimol.UniMolModel

Bases: Model

Multi-backend Uni-Mol model for 3D molecular representation learning and property prediction.

Supports molecular pretraining, candidate pocket pretraining, crystal, and OLED configurations.

Parameters:

Name Type Description Default
output_dim int

Number of task output dimensions / classes. (default: 2)

 2
data_type str

Data domain ("molecule", "protein", "crystal", "oled"). (default: "molecule")

 'molecule'
vocab_size int

Vocabulary size for token dictionary. (default: 512)

 512
encoder_layers int

Number of transformer encoder layers. (default: 15)

 15
encoder_embed_dim int

Node embedding dimension. (default: 512)

 512
encoder_ffn_embed_dim int

FFN hidden dimension. (default: 2048)

 2048
encoder_attention_heads int

Number of attention heads. (default: 64)

 64
kernel str

GBF kernel type ("gaussian" or "numerical"). (default: "gaussian")

 'gaussian'
num_kernel int

Number of radial kernels. (default: 128)

 128
pooler_dropout float

Dropout for classification head. (default: 0.0)

 0.0
activation_fn str

Activation function name. (default: "gelu")

 'gelu'
post_ln bool

Post-LN flag. (default: False)

 False
**kwargs

Additional model arguments.

 {}
Example 
import numpy as np
from k3_node.models import UniMolModel

tokens = np.random.randint(1, 64, size=(2, 6))  # atom tokens of 2 molecules with 6 atoms
coords = np.random.rand(2, 6, 3).astype("float32") * 3.0  # 3D conformations

model = UniMolModel(output_dim=2, vocab_size=64, encoder_layers=2, encoder_embed_dim=32, encoder_ffn_embed_dim=64,
                    encoder_attention_heads=4, num_kernel=16)
logits = model(tokens, src_coord=coords)  # molecule-level predictions
print(tuple(logits.shape))  # (2, 2)
reprs = model(tokens, src_coord=coords, return_repr=True)
print(tuple(reprs["cls_repr"].shape), tuple(reprs["encoder_rep"].shape))  # (2, 32) (2, 6, 32): molecule and atom embeddings

UniMolConfGenModel

A UniMolModel variant specialized for 3D conformation generation (denoising a noisy input conformation back to equilibrium geometry).

k3_node.models.unimol.UniMolConfGenModel

Bases: UniMolModel

Uni-Mol Conformation Generation Model for iterative 3D geometry prediction.

Example 
import numpy as np
from k3_node.models import UniMolConfGenModel

tokens = np.random.randint(1, 64, size=(2, 6))  # atom tokens of 2 molecules with 6 atoms
coords = np.random.rand(2, 6, 3).astype("float32") * 3.0  # 3D conformations

model = UniMolConfGenModel(vocab_size=64, encoder_layers=2, encoder_embed_dim=32, encoder_ffn_embed_dim=64,
                           encoder_attention_heads=4, num_kernel=16)
new_coords, pred_dist = model(tokens, src_coord=coords)  # refined conformation, predicted distances
print(tuple(new_coords.shape), tuple(pred_dist.shape))  # (2, 6, 3) (2, 6, 6)

Checkpoint Utilities

k3_node.models.unimol.load_unimol_weights(model, checkpoint_path=None, pretrained_name=None, folder='checkpoints', download=True)

Loads pre-trained weights from a PyTorch checkpoint (.pt) into a Keras 3 UniMolModel.

Parameters:

Name Type Description Default
model UniMolModel

The target UniMolModel instance.

 required
checkpoint_path str

Local path to .pt file or checkpoint name.

 None
pretrained_name str

Pretrained model identifier.

 None
folder str

Directory to store downloaded checkpoints. (default: "checkpoints")

 'checkpoints'
download bool

Whether to download checkpoint if missing locally. (default: True)

 True

Returns:

Name Type Description
UniMolModel UniMolModel

The model with loaded weights.

k3_node.models.unimol.download_unimol_checkpoint(name='mol_pre_no_h', folder='checkpoints', log=True)

Downloads a pre-trained Uni-Mol checkpoint (.pt).

Parameters:

Name Type Description Default
name str

Pretrained checkpoint name or alias (e.g. "mol_pre_no_h", "mol_pre_all_h", "pocket_pre", "mp_all_h", "oled_pre_no_h", "qm9", "drugs", "binding_pose").

 'mol_pre_no_h'
folder str

Target directory to save the checkpoint. (default: "checkpoints")

 'checkpoints'
log bool

Whether to print download progress. (default: True)

 True

Returns:

Name Type Description
str str

Absolute path to the downloaded checkpoint file.

Usage Example

from k3_node.models.unimol import UniMolModel, load_unimol_weights, download_unimol_checkpoint

ckpt_path = download_unimol_checkpoint("mol_pre_no_h")

model = UniMolModel(output_dim=2, data_type="molecule")
load_unimol_weights(model, checkpoint_path=ckpt_path)

UniMol2Model

Scalable Uni-Mol2 pretraining model, available in five parameter-count presets ("84m" through "1.1b").

k3_node.models.unimol2.UniMol2Model

Bases: Model

Scalable molecular pretraining model Uni-Mol2 (84M to 1.1B parameters).

Parameters:

Name Type Description Default
model_size str

Model scale preset ("84m", "164m", "310m", "570m", or "1.1b"). (default: "84m")

 '84m'
output_dim int

Classification/regression head dimension. (default: 2)

 2
num_encoder_layers int

Number of two-track transformer layers.

 None
encoder_embed_dim int

Atom feature dimension.

 None
num_attention_heads int

Number of attention heads.

 None
pair_embed_dim int

Pair feature dimension.

 None
**kwargs

Additional model arguments.

 {}
Example 
import numpy as np
from k3_node.models import UniMol2Model

atom_types = np.random.randint(0, 128, size=(2, 6))
coords = np.random.rand(2, 6, 3).astype("float32") * 3.0
model = UniMol2Model(output_dim=2, num_encoder_layers=2, encoder_embed_dim=32, num_attention_heads=4,
                     pair_embed_dim=16, ffn_embedding_dim=64, pair_hidden_dim=8)
print(tuple(model(atom_types, coords=coords).shape))  # (2, 2)
logits, new_coords, x, pair = model(atom_types, coords=coords, return_coords=True, return_repr=True)
print(tuple(new_coords.shape), tuple(x.shape), tuple(pair.shape))  # (2, 6, 3) (2, 6, 32) (2, 6, 6, 16)

Checkpoint Utilities

k3_node.models.unimol2.load_unimol2_weights(model, checkpoint_path=None, model_size=None, folder='checkpoints', download=True)

Loads weights from a PyTorch checkpoint into a Keras 3 UniMol2Model.

k3_node.models.unimol2.download_unimol2_checkpoint(model_size='84m', folder='checkpoints', log=True)

Downloads a pre-trained Uni-Mol2 checkpoint (.pt).

Parameters:

Name Type Description Default
model_size str

Model size ("84m", "164m", "310m", "570m", "1.1b").

 '84m'
folder str

Target folder. (default: "checkpoints")

 'checkpoints'
log bool

Print download progress. (default: True)

 True

Returns:

Name Type Description
str str

Absolute path to the downloaded file.

UniMolPlusPCQModel / UniMolPlusOC20Model

Uni-Mol+ two-track (atom + pair) transformer, pretrained on the PCQM4Mv2 quantum-chemistry benchmark (UniMolPlusPCQModel) or the OC20 catalyst dataset (UniMolPlusOC20Model).

k3_node.models.unimol_plus.UniMolPlusPCQModel

Bases: Model

Uni-Mol+ Model for quantum chemical property prediction on PCQM4Mv2.

Parameters:

Name Type Description Default
num_layers int

Number of iterative refinement layers. (default: 12)

 12
embed_dim int

Atom representation dimension. (default: 768)

 768
pair_dim int

Pair representation dimension. (default: 512)

 512
num_heads int

Number of attention heads. (default: 48)

 48
output_dim int

Number of predicted quantum properties. (default: 1)

 1
**kwargs

Additional model arguments.

 {}
Example 
import numpy as np
from k3_node.models import UniMolPlusPCQModel

atom_types = np.random.randint(0, 64, size=(2, 5))
coords = np.random.rand(2, 5, 3).astype("float32") * 3.0  # initial (e.g. RDKit) conformation
model = UniMolPlusPCQModel(num_layers=2, embed_dim=32, pair_dim=16, num_heads=4, output_dim=1)
pred = model(atom_types, coords=coords)  # HOMO-LUMO gap (PCQM4Mv2) per structure
print(tuple(pred.shape))  # (2, 1)

k3_node.models.unimol_plus.UniMolPlusOC20Model

Bases: UniMolPlusPCQModel

Uni-Mol+ Model for initial structure to relaxed energy prediction on OC20.

Example 
import numpy as np
from k3_node.models import UniMolPlusOC20Model

atom_types = np.random.randint(0, 64, size=(2, 5))
coords = np.random.rand(2, 5, 3).astype("float32") * 3.0  # initial (e.g. RDKit) conformation
model = UniMolPlusOC20Model(num_layers=2, embed_dim=32, pair_dim=16, num_heads=4, output_dim=1)
pred = model(atom_types, coords=coords)  # relaxed energy (OC20) per structure
print(tuple(pred.shape))  # (2, 1)

Checkpoint Utilities

k3_node.models.unimol_plus.load_unimol_plus_weights(model, checkpoint_path=None, pretrained_name=None, folder='checkpoints', download=True)

Loads pre-trained weights into a Keras 3 Uni-Mol+ model.

k3_node.models.unimol_plus.download_unimol_plus_checkpoint(name='pcq_base', folder='checkpoints', log=True)

Downloads a pre-trained Uni-Mol+ checkpoint (.pt).

Parameters:

Name Type Description Default
name str

Model name or alias ("pcq_base", "pcq_large", "pcq_small", "oc20_base").

 'pcq_base'
folder str

Target download directory. (default: "checkpoints")

 'checkpoints'
log bool

Print download log. (default: True)

 True

Returns:

Name Type Description
str str

Absolute path to the downloaded file.