| """ |
| Sample from a trained NanoDiffusionGPT checkpoint. |
| |
| Unlike GPT sampling, this does not append one token at a time. It creates a |
| fixed number of [MASK] tokens after the prompt and denoises them in parallel. |
| """ |
|
|
| import os |
| import pickle |
| from contextlib import nullcontext |
|
|
| import torch |
|
|
| from model import NanoDiffusionGPT, NanoDiffusionGPTConfig |
|
|
| |
| out_dir = "out-diffusion" |
| start = "\n" |
| num_samples = 5 |
| max_new_tokens = 300 |
| steps = 128 |
| temperature = 0.8 |
| top_k = 20 |
| seed = 1337 |
| device = "cuda" |
| dtype = "bfloat16" if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else "float16" |
| compile = False |
| if os.path.exists("configurator.py"): |
| exec(open("configurator.py").read()) |
| |
|
|
| torch.manual_seed(seed) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed(seed) |
| torch.backends.cuda.matmul.allow_tf32 = True |
| torch.backends.cudnn.allow_tf32 = True |
| device_type = "cuda" if "cuda" in device else "cpu" |
| ptdtype = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[dtype] |
| ctx = nullcontext() if device_type == "cpu" else torch.amp.autocast(device_type=device_type, dtype=ptdtype) |
|
|
| ckpt_path = os.path.join(out_dir, "ckpt.pt") |
| checkpoint = torch.load(ckpt_path, map_location=device) |
| model = NanoDiffusionGPT(NanoDiffusionGPTConfig(**checkpoint["model_args"])) |
| state_dict = checkpoint["model"] |
| unwanted_prefix = "_orig_mod." |
| for k, v in list(state_dict.items()): |
| if k.startswith(unwanted_prefix): |
| state_dict[k[len(unwanted_prefix) :]] = state_dict.pop(k) |
| model.load_state_dict(state_dict) |
| model.eval() |
| model.to(device) |
| if compile: |
| model = torch.compile(model) |
| raw_model = model._orig_mod if hasattr(model, "_orig_mod") else model |
|
|
| dataset = checkpoint["config"]["dataset"] |
| meta_path = os.path.join("data", dataset, "meta.pkl") |
| if not os.path.exists(meta_path): |
| meta_path = os.path.join(out_dir, "meta.pkl") |
| if not os.path.exists(meta_path): |
| meta_path = "meta.pkl" |
| print(f"Loading meta from {meta_path}...") |
| with open(meta_path, "rb") as f: |
| meta = pickle.load(f) |
| stoi, itos = meta["stoi"], meta["itos"] |
|
|
|
|
| def encode(s): |
| return [stoi[c] for c in s if c in stoi] |
|
|
|
|
| def decode(ids): |
| return "".join(itos[i] for i in ids if i in itos) |
|
|
|
|
| if start.startswith("FILE:"): |
| with open(start[5:], "r", encoding="utf-8") as f: |
| start = f.read() |
| start_ids = encode(start) |
| if not start_ids: |
| raise ValueError("prompt has no characters that exist in the dataset vocabulary") |
| if len(start_ids) + max_new_tokens > raw_model.config.block_size: |
| raise ValueError( |
| f"prompt length ({len(start_ids)}) + max_new_tokens ({max_new_tokens}) " |
| f"exceeds block_size ({raw_model.config.block_size})" |
| ) |
| x = torch.tensor(start_ids, dtype=torch.long, device=device)[None, ...] |
|
|
| with torch.no_grad(): |
| with ctx: |
| for k in range(num_samples): |
| y = model.generate( |
| x, |
| max_new_tokens=max_new_tokens, |
| steps=steps, |
| temperature=temperature, |
| top_k=top_k, |
| ) |
| print(decode(y[0].tolist())) |
| print("---------------") |
|
|
