| #!/usr/bin/env python3 | |
| """Minimal loader for the Cubic Hier 150M checkpoint.""" | |
| import json | |
| from pathlib import Path | |
| import torch | |
| from safetensors.torch import load_file | |
| from tokenizers import Tokenizer | |
| import modeling_cubic as M | |
| HERE = Path(__file__).resolve().parent | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| config_data = json.loads((HERE / "config.json").read_text(encoding="utf-8")) | |
| fields = {k: v for k, v in config_data.items() if k in M.Config.__dataclass_fields__} | |
| config = M.Config(**fields) | |
| model = M.CubicHierLM(config).to(DEVICE, dtype=torch.bfloat16).eval() | |
| model.load_state_dict(load_file(str(HERE / "model.safetensors"), device="cpu"), strict=True) | |
| model.gradient_checkpointing = False | |
| tokenizer = Tokenizer.from_file(str(HERE / "tokenizer.json")) | |
| print(M.generate(model, tokenizer, "Explain why the sky is blue.", max_new_tokens=160)) | |