| import argparse |
| import torch |
| from tokenizers import Tokenizer |
|
|
| from model import GPT, sample |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("ckpt", help="чекпоинт: kit_*.pt или last.pt") |
| ap.add_argument("prompt", nargs="?", default="Москва — это") |
| ap.add_argument("--tokens", type=int, default=150) |
| ap.add_argument("--temp", type=float, default=0.8) |
| ap.add_argument("--top_k", type=int, default=40) |
| args = ap.parse_args() |
|
|
| device = "cuda" |
| tok = Tokenizer.from_file("data/tokenizer.json") |
| model = GPT().to(device) |
| if args.ckpt.endswith(".safetensors"): |
| from safetensors.torch import load_file |
| |
| model.load_state_dict({k: v.float() for k, v in load_file(args.ckpt).items()}, strict=False) |
| else: |
| sd = torch.load(args.ckpt, map_location=device) |
| model.load_state_dict(sd["model"] if "model" in sd else sd) |
| print(sample(model, tok, args.prompt, device, args.tokens, args.temp, args.top_k)) |
|
|
