import numpy as np
from pathlib import Path

path = Path("model-batch.npz")
features = np.arange(12, dtype=np.float32).reshape(3, 4)
labels = np.array([0, 1, 1], dtype=np.int64)
sample_ids = np.array(["train-001", "train-002", "train-003"])

np.savez_compressed(
    path,
    features=features,
    labels=labels,
    sample_ids=sample_ids,
)

with np.load(path) as data:
    print("keys:", sorted(data.files))
    print("features shape:", data["features"].shape)
    print("features dtype:", data["features"].dtype)
    print("labels:", data["labels"].tolist())
    print("sample ids:", data["sample_ids"].tolist())
    print("features match:", np.array_equal(features, data["features"]))
    print("labels match:", np.array_equal(labels, data["labels"]))

print("file exists:", path.exists())