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