laya_head.py

"""Load the Laya decision head + config from <name>-head.safetensors (sits next to the GGUF backbone)."""
import json, sys
from safetensors import safe_open

def load_head(path):
    """-> (config dict, {tensor_name: torch.float32 tensor}); names match the original PyTorch state_dict."""
    with safe_open(path, "pt") as f:
        cfg = json.loads(f.metadata()["laya.config"])
        return cfg, {k: f.get_tensor(k) for k in f.keys()}

if __name__ == "__main__":
    cfg, h = load_head(sys.argv[1]); print(len(h), "head tensors;", list(cfg)[:5])