| """Load snapshots written by scripts/train.py (snap_*.pt / final.pt), full checkpoints, or an exported |
| directory (model.safetensors + config.json, as published on Hugging Face).""" |
| from __future__ import annotations |
|
|
| import json |
| import os |
|
|
| import torch |
|
|
| from tiny_agent.model import ModelConfig, TinyAgentLM |
|
|
|
|
| def load_model(path: str, device="xpu", dtype=torch.float32) -> TinyAgentLM: |
| if os.path.isdir(path): |
| from safetensors.torch import load_file |
| st = {"model": load_file(os.path.join(path, "model.safetensors")), |
| "config": json.load(open(os.path.join(path, "config.json")))} |
| else: |
| st = torch.load(path, map_location="cpu", weights_only=False) |
| cfg_d = st.get("config") or st["meta"]["config"] |
| for k in ("engram_layers", "engram_orders"): |
| cfg_d[k] = tuple(cfg_d[k]) |
| model = TinyAgentLM(ModelConfig(**cfg_d)) |
| model.load_state_dict({k: v.float() if v.is_floating_point() else v for k, v in st["model"].items()}) |
| return model.to(device=device, dtype=dtype) |
|
|
|
|
| def export(path: str, out_dir: str) -> None: |
| """Write a checkpoint as out_dir/model.safetensors (bf16) + out_dir/config.json.""" |
| from safetensors.torch import save_file |
| st = torch.load(path, map_location="cpu", weights_only=False) |
| cfg_d = dict(st.get("config") or st["meta"]["config"]) |
| sd = {k: (v.to(torch.bfloat16) if v.is_floating_point() else v).contiguous() for k, v in st["model"].items()} |
| os.makedirs(out_dir, exist_ok=True) |
| save_file(sd, os.path.join(out_dir, "model.safetensors")) |
| json.dump({k: list(v) if isinstance(v, tuple) else v for k, v in cfg_d.items()}, |
| open(os.path.join(out_dir, "config.json"), "w"), indent=1) |
|
|
