AstrAI/astrai/serialization/__init__.py

44 lines
860 B
Python

"""Serialization utilities for models and datasets.
This package re-exports checkpoint helpers and dataset storage helpers so
that existing imports from ``astrai.serialization`` continue to work.
"""
from astrai.serialization.checkpoint import (
Checkpoint,
load_json,
load_model_config,
load_model_weights,
load_safetensors,
load_state_dict,
load_torch,
save_json,
save_model,
save_safetensors,
save_torch,
)
from astrai.serialization.dataset import (
load_bin,
load_h5,
save_bin,
save_h5,
)
__all__ = [
"Checkpoint",
"load_json",
"load_model_config",
"load_model_weights",
"load_safetensors",
"load_state_dict",
"load_torch",
"save_json",
"save_model",
"save_safetensors",
"save_torch",
"load_bin",
"load_h5",
"save_bin",
"save_h5",
]