"""Minimal loader for Stasis 1. Run from a clone/snapshot of this repo.
pip install torch safetensors tokenizers huggingface_hub"""
import json, sys, torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
from tokenizers import Tokenizer

path = snapshot_download("ShivankXD/stasis-1", allow_patterns=["config.json", "tokenizer.json", "model.safetensors", "frozen_stasis/*"])  # skip checkpoints/ (large)
sys.path.insert(0, path)
from frozen_stasis.config import ModelConfig
from frozen_stasis.model import StasisDecoder

cfg = json.load(open(f"{path}/config.json"))
keys = ("vocab_size", "layers", "hidden_size", "heads", "intermediate_size", "context_length", "rope_theta", "rms_epsilon")
model = StasisDecoder(ModelConfig(**{k: cfg[k] for k in keys}))
model.load_state_dict(load_file(f"{path}/model.safetensors"), strict=True)
model.eval()
tok = Tokenizer.from_file(f"{path}/tokenizer.json")

prompt = "The water cycle describes how water moves through the environment."
ids = torch.tensor([tok.encode(prompt).ids])
out = model.generate(ids, max_new_tokens=60, temperature=0.8, top_k=50)
print(tok.decode(out[0].tolist()))