File size: 1,010 Bytes
8c967ea | 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 | 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
# head.weight связан с tok_emb.weight и в файле отсутствует
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))
|