File size: 2,348 Bytes
4719692 | 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 | """Generate text with the standalone Python model with the fixed evaluation sampling policy."""
import argparse
import json
from pathlib import Path
import torch
from model_100m.model import GPT, ModelConfig
from safetensors.torch import load_file
from tokenizers import Tokenizer
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--prompt", default="The solar system consists of")
p.add_argument("--tokens", type=int, default=128)
p.add_argument("--seed", type=int, default=42)
p.add_argument("--temperature", type=float, default=0.8, help="0 selects greedy decoding")
p.add_argument("--top-k", type=int, default=50, help="0 samples the full vocabulary")
p.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
a = p.parse_args()
if a.tokens < 0 or a.temperature < 0 or a.top_k < 0:
p.error("Generation settings must be nonnegative")
root = Path(__file__).resolve().parent
torch.set_num_threads(2)
torch.manual_seed(a.seed)
model = GPT(ModelConfig(**json.loads((root / "config.json").read_text()))).to(a.device).eval()
model.load_state_dict(load_file(str(root / "model.safetensors")), strict=True)
tokenizer = Tokenizer.from_file(str(root / "tokenizer/tokenizer.json"))
ids = [1] + tokenizer.encode(a.prompt, add_special_tokens=False).ids
generator = torch.Generator(device=a.device).manual_seed(a.seed)
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
with torch.inference_mode():
for _ in range(a.tokens):
logits = model(torch.tensor([ids[-model.config.block_size :]], device=a.device))[
0, -1
].float()
if not torch.isfinite(logits).all():
raise ValueError("Non-finite generation logits")
logits[[0, 1, 3]] = -torch.inf # PAD/BOS/UNK; EOS remains eligible.
if a.temperature == 0:
token = int(logits.argmax())
elif a.top_k:
values, indices = torch.topk(logits / a.temperature, min(a.top_k, logits.numel()))
token = int(indices[torch.multinomial(values.softmax(-1), 1, generator=generator)])
else:
token = int(
torch.multinomial((logits / a.temperature).softmax(-1), 1, generator=generator)
)
ids.append(token)
if token == 2:
break
print(tokenizer.decode(ids, skip_special_tokens=True))
|