| import argparse |
| from pathlib import Path |
|
|
| import torch |
|
|
| from train import TinyTransformerLM |
|
|
|
|
| @torch.no_grad() |
| def generate(model, idx, max_new_tokens, temperature, itos): |
| model.eval() |
| for _ in range(max_new_tokens): |
| idx_cond = idx[:, -model.block_size :] |
| logits, _ = model(idx_cond) |
| logits = logits[:, -1, :] / temperature |
| probs = torch.softmax(logits, dim=-1) |
| next_id = torch.multinomial(probs, num_samples=1) |
| idx = torch.cat((idx, next_id), dim=1) |
| return "".join(itos[int(i)] for i in idx[0]) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model", default="runs/tiny-char-model.pt") |
| parser.add_argument("--prompt", default="hello") |
| parser.add_argument("--tokens", type=int, default=400) |
| parser.add_argument("--temperature", type=float, default=0.8) |
| args = parser.parse_args() |
|
|
| checkpoint = torch.load(Path(args.model), map_location="cpu") |
| config = checkpoint["config"] |
| stoi = checkpoint["stoi"] |
| itos = {int(k): v for k, v in checkpoint["itos"].items()} |
|
|
| model = TinyTransformerLM(**config) |
| model.load_state_dict(checkpoint["model"]) |
|
|
| fallback = next(iter(stoi.values())) |
| encoded = [stoi.get(ch, fallback) for ch in args.prompt] |
| idx = torch.tensor([encoded], dtype=torch.long) |
| print(generate(model, idx, args.tokens, args.temperature, itos)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
