hyperglyph-codec 0.1.0__py3-none-any.whl

hyperglyph/serialization.py

@@ -0,0 +1,75 @@
+"""Serialization helpers for compressed models."""
+
+from __future__ import annotations
+
+import json
+import zipfile
+from pathlib import Path
+from typing import Any, Mapping
+
+import numpy as np
+
+from .codec import CompressedModel, CompressedTensor
+
+
+def save_compressed(compressed_model: CompressedModel, path: str | Path) -> None:
+    """Save a compressed model to a .hwz zip archive."""
+    destination = Path(path)
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    with zipfile.ZipFile(destination, "w", compression=zipfile.ZIP_DEFLATED) as archive:
+        metadata = {
+            "format_version": compressed_model.format_version,
+            "tensors": {
+                name: tensor_to_dict(tensor) for name, tensor in compressed_model.tensors.items()
+            },
+        }
+        archive.writestr("metadata.json", json.dumps(metadata, indent=2))
+        prototype_arrays = {}
+        for name, tensor in compressed_model.tensors.items():
+            prototype_arrays[f"{name}_prototypes"] = tensor.prototype_matrix
+        if prototype_arrays:
+            with archive.open("prototypes.npz", "w") as handle:
+                np.savez(handle, **prototype_arrays)  # type: ignore[arg-type]
+
+
+def load_compressed(path: str | Path) -> CompressedModel:
+    """Load a compressed model from a .hwz archive."""
+    archive_path = Path(path)
+    with zipfile.ZipFile(archive_path, "r") as archive:
+        metadata = json.loads(archive.read("metadata.json"))
+        tensors: dict[str, CompressedTensor] = {}
+        for name, value in metadata.get("tensors", {}).items():
+            tensors[name] = dict_to_tensor(value)
+        return CompressedModel(
+            tensors=tensors, payload=b"", format_version=metadata.get("format_version", "0.1")
+        )
+
+
+def tensor_to_dict(tensor: CompressedTensor) -> dict[str, Any]:
+    """Convert a compressed tensor to a JSON-safe dictionary."""
+    return {
+        "name": tensor.name,
+        "shape": list(tensor.shape),
+        "block_size": tensor.block_size,
+        "prototype_ids": tensor.prototype_ids,
+        "scales": tensor.scales,
+        "residuals": tensor.residuals,
+        "prototype_matrix": tensor.prototype_matrix.tolist(),
+        "seed": tensor.seed,
+        "codec_config": tensor.codec_config,
+    }
+
+
+def dict_to_tensor(payload: Mapping[str, Any]) -> CompressedTensor:
+    """Convert a JSON-safe dictionary back to a CompressedTensor."""
+    return CompressedTensor(
+        name=str(payload["name"]),
+        shape=tuple(int(value) for value in payload["shape"]),
+        block_size=int(payload["block_size"]),
+        prototype_ids=[int(value) for value in payload["prototype_ids"]],
+        scales=[float(value) for value in payload["scales"]],
+        residuals=[dict(value) for value in payload["residuals"]],
+        prototype_matrix=np.asarray(payload["prototype_matrix"], dtype=np.float32),
+        seed=int(payload["seed"]),
+        codec_config=dict(payload["codec_config"]),
+    )

hyperglyph/torch_adapter.py

@@ -0,0 +1,67 @@
+"""Optional PyTorch adapter."""
+
+from __future__ import annotations
+
+from typing import Any, Mapping
+
+from .config import HyperGlyphConfig
+from .exceptions import OptionalDependencyError
+
+
+def is_torch_available() -> bool:
+    """Check whether torch is installed."""
+    try:
+        import torch  # noqa: F401
+    except ImportError:
+        return False
+    return True
+
+
+def tensor_to_numpy(tensor: Any) -> Any:
+    """Convert a torch tensor to a NumPy array."""
+    if not is_torch_available():
+        raise OptionalDependencyError("torch is required for tensor_to_numpy")
+    import torch
+
+    if isinstance(tensor, torch.Tensor):
+        return tensor.detach().cpu().numpy()
+    return tensor
+
+
+def numpy_to_tensor(array: Any, reference_tensor: Any | None = None) -> Any:
+    """Convert a NumPy array back to a torch tensor."""
+    if not is_torch_available():
+        raise OptionalDependencyError("torch is required for numpy_to_tensor")
+    import torch
+
+    if reference_tensor is not None and isinstance(reference_tensor, torch.Tensor):
+        return torch.tensor(array, dtype=reference_tensor.dtype, device=reference_tensor.device)
+    return torch.tensor(array)
+
+
+def compress_state_dict(
+    state_dict: Mapping[str, Any], config: HyperGlyphConfig | None = None
+) -> Any:
+    """Compress a torch state_dict."""
+    if not is_torch_available():
+        raise OptionalDependencyError("torch is required for compress_state_dict")
+    from .codec import HyperGlyphCodec
+
+    codec = HyperGlyphCodec(config)
+    numpy_state_dict = {name: tensor_to_numpy(value) for name, value in state_dict.items()}
+    return codec.compress_state_dict(numpy_state_dict)
+
+
+def decompress_state_dict(
+    compressed_model: Any, reference_state_dict: Mapping[str, Any] | None = None
+) -> dict[str, Any]:
+    """Decompress a compressed model into a dictionary of numpy arrays or torch tensors."""
+    if not is_torch_available():
+        raise OptionalDependencyError("torch is required for decompress_state_dict")
+    from .codec import HyperGlyphCodec
+
+    codec = HyperGlyphCodec()
+    restored = codec.decompress_state_dict(compressed_model)
+    if reference_state_dict is None:
+        return restored
+    return {name: numpy_to_tensor(restored[name], reference_state_dict[name]) for name in restored}