Hugging Face Hub API Reference

k3_node.hub.hub_mixin.K3NodeHubMixin

Mixin class providing seamless saving, loading, and publishing of K3-Node GNN models and tasks to/from the Hugging Face Hub.

Methods:

Name	Description
save_pretrained	Saves model weights, config, and Model Card to a local directory.
from_pretrained	Loads a model from a local folder or Hugging Face Hub repository.
push_to_hub	Automatically saves and pushes the model to a Hugging Face Hub repo.
predict	Runs inference on graph data (PyG Data, molecular structures, dicts, or tensors).

export_onnx(output_path, dummy_inputs=None, opset=17, dynamic_axes=True, **kwargs)

Exports this model or task to high-performance ONNX format.

Parameters:

Name	Type	Description	Default
output_path	Union[str, Path]	Target path for the .onnx file.	required
dummy_inputs	Optional[Any]	Optional sample input data.	None
opset	int	ONNX operator set version. (default: 17)	17
dynamic_axes	bool	Whether graph size dimensions are dynamic. (default: True)	True

Returns:

Type	Description
Path	Path object pointing to the generated .onnx file.

export_tensorrt(output_path, dummy_inputs=None, precision='fp16', **kwargs)

Compiles this model or task into a high-throughput NVIDIA TensorRT engine.

Parameters:

Name	Type	Description	Default
output_path	Union[str, Path]	Target path for the .engine binary.	required
dummy_inputs	Optional[Any]	Optional sample input data.	None
precision	str	Precision mode ("fp32", "fp16", "int8").	'fp16'

Returns:

Type	Description
Path	Path object pointing to the generated TensorRT .engine file.

export_tflite(output_path, dummy_inputs=None, quantization=None, **kwargs)

Exports this model or task to an optimized TensorFlow Lite flatbuffer.

Parameters:

Name	Type	Description	Default
output_path	Union[str, Path]	Target path for the .tflite file.	required
dummy_inputs	Optional[Any]	Optional sample input data.	None
quantization	Optional[str]	Quantization mode (None, "fp16", "int8_dynamic", "int8_full").	None

Returns:

Type	Description
Path	Path object pointing to the generated .tflite file.

from_pretrained(repo_id_or_path, revision=None, token=None, cache_dir=None, **model_kwargs) classmethod

Loads a pretrained K3-Node task or model from a local folder or Hugging Face Hub.

Parameters:

Name	Type	Description	Default
repo_id_or_path	Union[str, Path]	Local directory path or Hugging Face repo ID (e.g. "k3-node/schnet-qm9").	required
revision	Optional[str]	Specific git revision/branch on Hugging Face Hub.	None
token	Optional[Union[str, bool]]	Hugging Face authentication token.	None
cache_dir	Optional[Union[str, Path]]	Cache directory for downloaded Hub files.	None
**model_kwargs		Overrides for configuration parameters.	{}

Returns:

Type	Description
T	Restored and initialized model or task instance with loaded weights.

predict(data=None, *args, **kwargs)

Infers predictions on graph or molecular data.

Supports PyG / K3-Node Data objects (extracting (z, pos, batch) for molecular models or (x, edge_index, ...) for standard GNNs), dictionaries, tuples of tensors, or direct positional tensors.

Parameters:

Name	Type	Description	Default
data	Any	Input graph or molecule Data, dict, or tensor.	None
*args	Any	Additional positional arguments.	()
**kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
Any	Model prediction tensor or array.

push_to_hub(repo_id, token=None, private=False, commit_message=None, metrics=None, dataset_name=None, license='mit', **kwargs)

Saves the model and pushes it directly to the Hugging Face Hub.

Parameters:

Name	Type	Description	Default
repo_id	str	Hugging Face repo ID in format "username/model_name" or "org/model_name".	required
token	Optional[Union[str, bool]]	Optional Hugging Face auth token. If not passed, uses cached credentials.	None
private	bool	Whether the repository should be private. (default: False)	False
commit_message	Optional[str]	Optional commit message for the upload.	None
metrics	Optional[Dict[str, float]]	Optional dictionary of evaluation metrics to document in the Model Card.	None
dataset_name	Optional[str]	Optional dataset name for the Model Card.	None
license	str	License tag. (default: "mit")	'mit'

Returns:

Type	Description
str	Web URL of the repository on Hugging Face Hub.

save_pretrained(save_directory, config=None, metrics=None, dataset_name=None, repo_id=None, license='mit', **kwargs)

Saves model weights, config.json, and README.md (Model Card) to disk.

Parameters:

Name	Type	Description	Default
save_directory	Union[str, Path]	Directory path to save model files in.	required
config	Optional[Dict[str, Any]]	Optional custom configuration dictionary.	None
metrics	Optional[Dict[str, float]]	Optional evaluation metrics dictionary to include in Model Card.	None
dataset_name	Optional[str]	Optional dataset name for the Model Card.	None
repo_id	Optional[str]	Optional Hugging Face repository ID.	None
license	str	License identifier. (default: "mit")	'mit'

Returns:

Type	Description
Path	Path object of the saved directory.

k3_node.hub.hub_mixin.from_pretrained(repo_id_or_path, task_cls=None, **kwargs)

Loads a model or task from a local directory or Hugging Face Hub.

k3_node.hub.hub_mixin.save_pretrained(model_or_task, save_directory, **kwargs)

Saves a model or task to disk in Hugging Face Hub format.

k3_node.hub.hub_mixin.push_to_hub(model_or_task, repo_id, **kwargs)

Pushes a model or task directly to the Hugging Face Hub.

k3_node.hub.model_card.generate_model_card(task_type, backbone, config, metrics=None, dataset_name=None, repo_id=None, license='mit')

Generates a standard Hugging Face Model Card with YAML frontmatter and markdown documentation for a K3-Node GNN model.

Parameters:

Name	Type	Description	Default
task_type	str	Name of the task or model (e.g. 'NodeClassifier', 'SchNet').	required
backbone	str	Name of the backbone architecture (e.g. 'gcn', 'schnet', 'chgnet').	required
config	Dict[str, Any]	Dictionary containing model architecture and training hyperparameters.	required
metrics	Optional[Dict[str, float]]	Optional dictionary of evaluation metrics (e.g. {'accuracy': 0.82}).	None
dataset_name	Optional[str]	Optional name of the dataset the model was trained on.	None
repo_id	Optional[str]	Optional repository ID on Hugging Face Hub.	None
license	str	Open-source license tag. (default: 'mit')	'mit'

Returns:

Type	Description
str	Formatted markdown string representing the README.md model card.

k3_node.hub.dataset_hub.save_graph_dataset(dataset, filepath)

Saves a single graph (Data) or a list of graphs (List[Data]) into a compressed multi-backend .npz file.

Parameters:

Name	Type	Description	Default
dataset	Union[Data, List[Data]]	A :class:k3_node.data.Data object or list of :class:Data objects.	required
filepath	Union[str, Path]	Output .npz file path.	required

Returns:

Type	Description
Path	Path of the saved file.

k3_node.hub.dataset_hub.load_graph_dataset(filepath)

Loads a graph or collection of graphs from a .npz file.

Parameters:

Name	Type	Description	Default
filepath	Union[str, Path]	Path to the .npz dataset file.	required

Returns:

Name	Type	Description
A	Union[Data, List[Data]]	class:Data object or a list of :class:Data objects.

k3_node.hub.dataset_hub.push_dataset_to_hub(dataset, repo_id, token=None, private=False, commit_message=None, description=None, license='mit')

Pushes a graph dataset to the Hugging Face Hub under a dataset repository.

Parameters:

Name	Type	Description	Default
dataset	Union[Data, List[Data]]	A :class:k3_node.data.Data object or list of :class:Data objects.	required
repo_id	str	Hugging Face dataset repository ID (e.g. "username/my-graph-data").	required
token	Optional[Union[str, bool]]	Optional authentication token.	None
private	bool	Whether the repository should be private. (default: False)	False
commit_message	Optional[str]	Commit message for the upload.	None
description	Optional[str]	Optional dataset description.	None
license	str	License tag. (default: "mit")	'mit'

Returns:

Type	Description
str	URL of the dataset on Hugging Face Hub.

k3_node.hub.dataset_hub.load_dataset_from_hub(repo_id, filename='graph_data.npz', token=None, cache_dir=None)

Downloads and loads a graph dataset from the Hugging Face Hub.

Parameters:

Name	Type	Description	Default
repo_id	str	Hugging Face dataset repository ID (e.g. "username/my-graph-data").	required
filename	str	Name of the dataset file in the repo. (default: "graph_data.npz")	'graph_data.npz'
token	Optional[Union[str, bool]]	Optional authentication token.	None
cache_dir	Optional[Union[str, Path]]	Optional cache directory.	None

Returns:

Name	Type	Description
Restored	Union[Data, List[Data]]	class:Data or :class:List[Data] object.