| import torch, argparse, json |
| from tokenizers import Tokenizer |
| from model.tiny_gpt2 import TinyGPT2, GPTConfig |
|
|
| parser = argparse.ArgumentParser() |
| parser.add_argument("--prompt", type=str, required=True) |
| parser.add_argument("--ckpt", type=str, default="out/sft/model_sft.pt") |
| parser.add_argument("--cfg", type=str, default="out/pretrain/gpt_config.json") |
| parser.add_argument("--tok", type=str, default="out/tokenizer.json") |
| args = parser.parse_args() |
|
|
| tok = Tokenizer.from_file(args.tok) |
| cfg = GPTConfig(**json.load(open(args.cfg))) |
| m = TinyGPT2(cfg) |
| m.load_state_dict(torch.load(args.ckpt, map_location="cpu")) |
| m.eval() |
|
|
| ids = tok.encode("[BOS] " + args.prompt).ids |
| x = torch.tensor([ids], dtype=torch.long) |
| with torch.no_grad(): |
| y = m.generate(x, max_new_tokens=80) |
| text = tok.decode(y[0].tolist()) |
| print(text) |
|
|
