Skip to content

Hugging Face Hub Integration

K3-Node provides first-class, native integration with the Hugging Face Hub, allowing you to share, discover, and download Graph Neural Network models and graph datasets with a single line of code.

Because K3-Node is built on Keras 3, every model pushed to the Hub is truly multi-backend: you can upload a model trained with PyTorch and load it instantly for inference with JAX or TensorFlow, and vice versa!

Key Features

  • One-line Sharing & Loading: model.push_to_hub("username/cora-gcn") and NodeClassifier.from_pretrained("username/cora-gcn").
  • Automatic Model Cards: Automatically generates rich README.md model cards with metadata, graph-ml pipeline tags, hyperparameters, metrics, and copy-paste usage snippets.
  • Graph Dataset Hub: Push and load custom graph datasets (Data or List[Data]) directly to/from Hugging Face dataset repositories.
  • Universal Multi-Backend Compatibility: Saved weights and configs work across PyTorch, JAX, and TensorFlow backends.

1. Saving & Loading Models Locally

Before uploading to the Hub, you can save and load model checkpoints locally using the exact same standard Hugging Face format:

import os
os.environ["KERAS_BACKEND"] = "torch"  # or "jax", "tensorflow"

from k3_node.datasets import Planetoid
from k3_node.tasks import NodeClassifier

# Load Cora dataset and train a GCN classifier
dataset = Planetoid(root="/tmp/Cora", name="Cora")
data = dataset[0]

clf = NodeClassifier(backbone="gcn", hidden_channels=32, num_layers=2)
clf.fit(data, epochs=30)

# Evaluate
metrics = clf.evaluate(data, mask="test_mask")
print("Accuracy:", metrics["accuracy"])

# Save locally in Hugging Face Hub format
clf.save_pretrained("./my_cora_gcn", metrics=metrics, dataset_name="Cora")

This creates a standard directory structure: 

my_cora_gcn/
├── README.md               # Generated Model Card with tags and metrics
├── config.json             # Model architecture & hyperparameters
└── model.weights.h5        # Keras 3 neural network weights

Loading from Local Directory

# Load using the specific task class
loaded_clf = NodeClassifier.from_pretrained("./my_cora_gcn")

# Or load generically (K3-Node inspects config.json to auto-instantiate the right task!)
from k3_node.hub import from_pretrained
generic_clf = from_pretrained("./my_cora_gcn")

# Run predictions
preds = generic_clf.predict(data)

2. Authentication & Setup

Before creating or uploading models and datasets to the Hugging Face Hub, you need a Hugging Face user access token with write permissions:

  1. Create a Write Token: Generate a write-enabled access token at:
    👉 https://huggingface.co/settings/tokens/new?preset=write

  2. Log in to Hugging Face: In your terminal, run the official login command: 

    hf auth login

Once authenticated, your credentials will be cached locally and used automatically by push_to_hub(). You can also pass your token explicitly in Python using token="hf_...".

3. Publishing Models to Hugging Face Hub

Push your trained model directly to your Hugging Face account:

# Push to Hub (uses cached credentials from `hf auth login` or explicit token)
repo_url = clf.push_to_hub(
    repo_id="your-username/cora-node-gcn",
    metrics=metrics,
    dataset_name="Cora",
    commit_message="Initial release of trained Cora GCN",
    private=False,
    # token="hf_...",  # optional if logged in via `hf auth login`
)
print("Model published at:", repo_url)

4. Loading Pretrained Models from the Hub

Anyone can load and run your published model with a single line of code:

import os
os.environ["KERAS_BACKEND"] = "jax"  # Works across all backends!

from k3_node.tasks import NodeClassifier

# Load directly from the Hugging Face Hub
model = NodeClassifier.from_pretrained("your-username/cora-node-gcn")

# Run inference
predictions = model.predict(new_graph)
probabilities = model.predict_proba(new_graph)

All 4 high-level task estimators support save_pretrained, from_pretrained, and push_to_hub: - NodeClassifier - GraphClassifier - GraphRegressor - LinkPredictor

5. Direct Model Architecture Hub Integration (k3.models.*)

Beyond the task estimators, all models in k3_node.models natively support saving, loading, inference, and pushing to the Hugging Face Hub:

import k3_node as k3

# 1. Load pre-trained weights with one line
model = k3.models.SchNet.from_pretrained("k3-node/schnet-qm9")

# 2. Predict directly on molecular or graph Data
energy = model.predict(molecule_data)

# 3. Push community checkpoints directly to the hub
model.push_to_hub("anas-rz/chgnet-mp-2026")

Supported Model Families

Every model in k3_node.models has this capability built in, including: - 3D Molecular & Quantum GNNs: SchNet, DimeNet, DimeNetPlusPlus, ViSNet, GNNFF - Materials & Crystal Models: CHGNet, MEGNet, M3GNet, TensorNet, SO3Net - Classical & Modern GNNs: GCN, GraphSAGE, GIN, GAT, PNA, EdgeCNN - Transformers & Foundation Models: Graphormer, GPSModel, UniMolModel, MoleBERT

6. Graph Dataset Hub Integration

Sharing graph datasets (single graphs or collections of graphs) is just as simple:

Pushing a Graph Dataset to the Hub

from k3_node.hub import push_dataset_to_hub
from k3_node.datasets import TUDataset

# Load a collection of molecular graphs
mutag = TUDataset(root="/tmp/MUTAG", name="MUTAG")

# Push to Hugging Face Hub as a dataset repo
dataset_url = push_dataset_to_hub(
    dataset=list(mutag),
    repo_id="your-username/mutag-graphs",
    description="MUTAG mutagenic aromatic and heteroaromatic nitro compounds graph benchmark.",
)
print("Dataset published at:", dataset_url)

Loading a Graph Dataset from the Hub

from k3_node.hub import load_dataset_from_hub

# Load directly into K3-Node Data structures
graphs = load_dataset_from_hub("your-username/mutag-graphs")
print(f"Loaded {len(graphs)} graphs! Sample: {graphs[0]}")

7. Summary Table

Operation Model Hub Function / Method Dataset Hub Function
Save Locally model.save_pretrained("./dir") save_graph_dataset(data, "path.npz")
Load Locally Task.from_pretrained("./dir") load_graph_dataset("path.npz")
Push to Hub model.push_to_hub("org/repo") push_dataset_to_hub(data, "org/repo")
Load from Hub Task.from_pretrained("org/repo") load_dataset_from_hub("org/repo")