Asilarkness/cubic-150m

1#!/usr/bin/env python32"""Minimal loader for the Cubic Hier 150M checkpoint."""3 4import json5from pathlib import Path6 7import torch8from safetensors.torch import load_file9from tokenizers import Tokenizer10 11import modeling_cubic as M12 13HERE = Path(__file__).resolve().parent14DEVICE = "cuda" if torch.cuda.is_available() else "cpu"15 16config_data = json.loads((HERE / "config.json").read_text(encoding="utf-8"))17fields = {k: v for k, v in config_data.items() if k in M.Config.__dataclass_fields__}18config = M.Config(**fields)19 20model = M.CubicHierLM(config).to(DEVICE, dtype=torch.bfloat16).eval()21model.load_state_dict(load_file(str(HERE / "model.safetensors"), device="cpu"), strict=True)22model.gradient_checkpointing = False23tokenizer = Tokenizer.from_file(str(HERE / "tokenizer.json"))24 25print(M.generate(model, tokenizer, "Explain why the sky is blue.", max_new_tokens=160))26