Spaces:
Sleeping
Sleeping
File size: 2,514 Bytes
ba05e77 | 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | """Text generation from trained GPT checkpoint."""
import argparse
import torch
from model import GPT
@torch.no_grad()
def generate(
model: GPT,
prompt: str,
stoi: dict[str, int],
itos: dict[int, str],
max_new_tokens: int = 200,
temperature: float = 0.8,
top_k: int = 40,
) -> str:
"""Generate text from a prompt using the trained model."""
device = next(model.parameters()).device
tokens = [stoi[c] for c in prompt if c in stoi]
idx = torch.tensor([tokens], dtype=torch.long, device=device)
model.eval()
for _ in range(max_new_tokens):
idx_cond = idx[:, -model.config.block_size :]
logits, _ = model(idx_cond)
logits = logits[:, -1, :] / temperature
if top_k > 0:
values, _ = torch.topk(logits, top_k)
logits[logits < values[:, -1:]] = float("-inf")
probs = torch.softmax(logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
idx = torch.cat([idx, next_token], dim=1)
return "".join([itos[i] for i in idx[0].tolist()])
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Generate text from a trained GPT checkpoint"
)
parser.add_argument(
"checkpoint", help="Path to checkpoint file (e.g. checkpoint_final.pt)"
)
parser.add_argument(
"--prompt", default="கடவுள் வாழ்த்து", help="Starting text for generation"
)
parser.add_argument(
"--max_new_tokens", type=int, default=200, help="Number of tokens to generate"
)
parser.add_argument(
"--temperature",
type=float,
default=0.8,
help="Sampling temperature (lower = more deterministic)",
)
parser.add_argument(
"--top_k", type=int, default=40, help="Only sample from top-k most likely tokens"
)
parser.add_argument("--seed", type=int, default=None, help="Random seed")
args = parser.parse_args()
if args.seed is not None:
torch.manual_seed(args.seed)
checkpoint = torch.load(args.checkpoint, weights_only=False)
config = checkpoint["config"]
stoi = checkpoint["stoi"]
itos = checkpoint["itos"]
model = GPT(config)
model.load_state_dict(checkpoint["model_state_dict"])
output = generate(
model,
args.prompt,
stoi,
itos,
max_new_tokens=args.max_new_tokens,
temperature=args.temperature,
top_k=args.top_k,
)
print(output)
|