File size: 1,652 Bytes
5ccb4fd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | # Copyright (c) 2026 Simulacra Research Inc.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import json
from pathlib import Path
from typing import Any, Literal, TypeVar
import equinox as eqx
T = TypeVar("T")
CheckpointKind = Literal["router", "compiled_finetune"]
def _metadata_path(path: str | Path) -> Path:
source = Path(path)
return source.with_name(source.name + ".json")
def save_model(
path: str | Path,
model: object,
*,
kind: CheckpointKind | None = None,
metadata: dict[str, Any] | None = None,
) -> None:
destination = Path(path)
destination.parent.mkdir(parents=True, exist_ok=True)
eqx.tree_serialise_leaves(destination, model)
if kind is not None:
payload = {"kind": kind, **(metadata or {})}
_metadata_path(destination).write_text(
json.dumps(payload, indent=2, sort_keys=True) + "\n"
)
def load_model(path: str | Path, template: T) -> T:
return eqx.tree_deserialise_leaves(Path(path), template)
def load_model_metadata(path: str | Path) -> dict[str, Any] | None:
source = _metadata_path(path)
return json.loads(source.read_text()) if source.exists() else None
def save_mcmc(path: str | Path, state: object) -> None:
destination = Path(path)
destination.parent.mkdir(parents=True, exist_ok=True)
eqx.tree_serialise_leaves(destination, state)
def load_mcmc(path: str | Path, template: T) -> T:
return eqx.tree_deserialise_leaves(Path(path), template)
__all__ = [
"CheckpointKind",
"load_mcmc",
"load_model",
"load_model_metadata",
"save_mcmc",
"save_model",
]
|