Views
No views yet
tiktoken ("gpt2", 50257)1import os, sys, json, torch, tiktoken
2from huggingface_hub import hf_hub_download
3from safetensors.torch import load_file
4
5repo_id = "gopluto-ai/slm-124m-tinystories-wiki"
6
7cfg_path = hf_hub_download(repo_id, "config.json")
8wt_path = hf_hub_download(repo_id, "model.safetensors")
9slm_path = hf_hub_download(repo_id, "slm.py")
10
11# make the downloaded slm.py importable
12sys.path.append(os.path.dirname(slm_path))
13from slm import SLM
14
15# load model
16cfg = json.load(open(cfg_path))
17model = SLM(**cfg)
18model.load_state_dict(load_file(wt_path))
19model.eval()
20
21# tokeniser
22enc = tiktoken.get_encoding("gpt2")
23prompt = "Once upon a time"
24ids = torch.tensor(enc.encode(prompt)).unsqueeze(0)
25
26# generate
27with torch.no_grad():
28 y = model.generate(ids, max_new_tokens=50, temperature=0.9, top_k=50)
29print(enc.decode(y[0].tolist()))