#!/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))