File size: 2,938 Bytes
213807c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | """Complete a story with a locally trained Sprout checkpoint (not a chatbot)."""
import argparse
import contextlib
import hashlib
from pathlib import Path
import torch
from tokenizers import Tokenizer
from model import Sprout
@torch.inference_mode()
def generate(model, tokenizer, prompt, max_new_tokens=180, temperature=.8, top_k=40, seed=42):
if temperature <= 0 or top_k < 1 or max_new_tokens < 0:
raise ValueError("temperature and top_k must be positive; token count nonnegative")
device = next(model.parameters()).device
generator = torch.Generator(device=device).manual_seed(seed)
eot = tokenizer.token_to_id("<|endoftext|>")
ids = tokenizer.encode(prompt).ids or [eot]
x = torch.tensor([ids], device=device, dtype=torch.long)
was_training = model.training
model.eval()
try:
for _ in range(max_new_tokens):
ctx = torch.autocast("cuda", dtype=torch.bfloat16) if device.type == "cuda" else contextlib.nullcontext()
with ctx:
logits = model(x[:, -model.block_size:])[:, -1, :].float() / temperature
threshold = torch.topk(logits, min(top_k, logits.size(-1))).values[:, [-1]]
logits = logits.masked_fill(logits < threshold, -float("inf"))
token = torch.multinomial(torch.softmax(logits, dim=-1), 1, generator=generator)
if token.item() == eot:
break
x = torch.cat((x, token), dim=1)
return tokenizer.decode(x[0].tolist(), skip_special_tokens=True)
finally:
model.train(was_training)
def main():
root = Path(__file__).resolve().parent
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--checkpoint", type=Path, default=root / "runs/sprout/best.pt")
parser.add_argument("--tokenizer", type=Path, default=root / "data/tokenizer.json")
parser.add_argument("--prompt", default="Once upon a time, a tiny robot found a seed.")
parser.add_argument("--tokens", type=int, default=220)
parser.add_argument("--temperature", type=float, default=.8)
parser.add_argument("--top-k", type=int, default=40)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
args = parser.parse_args()
torch.set_num_threads(4)
saved = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
expected = saved.get("tokenizer_sha256")
if expected and hashlib.sha256(args.tokenizer.read_bytes()).hexdigest() != expected:
raise ValueError("Tokenizer does not match checkpoint")
model = Sprout(**saved["model_config"])
model.load_state_dict(saved["model"])
del saved
model.to(args.device)
tokenizer = Tokenizer.from_file(str(args.tokenizer))
print(generate(model, tokenizer, args.prompt, args.tokens, args.temperature, args.top_k, args.seed))
if __name__ == "__main__":
main()
|