| """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())) | |