File size: 1,933 Bytes
5b8e8b4 | 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 | """
Standalone inference for Sid-GPT.
pip install torch tokenizers
python generate.py --prompt "Responsible for the mobile app roadmap."
"""
import argparse
import os
import torch
from tokenizers import ByteLevelBPETokenizer
from model import GPT
HERE = os.path.dirname(os.path.abspath(__file__))
TEMPLATE = "Turn this into a strong, quantified PM resume bullet:\n{}"
def load(checkpoint=None, tokenizer_dir=None, device="cpu"):
tok = ByteLevelBPETokenizer(os.path.join(tokenizer_dir or os.path.join(HERE, "tokenizer"), "vocab.json"),
os.path.join(tokenizer_dir or os.path.join(HERE, "tokenizer"), "merges.txt"))
ckpt = torch.load(checkpoint or os.path.join(HERE, "sft.pt"), map_location=device, weights_only=False)
model = GPT(ckpt["config"]).to(device)
model.load_state_dict(ckpt["model"])
model.eval()
return model, tok
def rewrite(model, tok, bullet, temperature=0.5, top_k=20, max_new_tokens=70, device="cpu"):
uid, aid, eot = (tok.token_to_id(t) for t in ("<|user|>", "<|assistant|>", "<|endoftext|>"))
ids = [uid] + tok.encode(TEMPLATE.format(bullet.strip())).ids + [aid]
with torch.no_grad():
out = model.generate(torch.tensor([ids], device=device), max_new_tokens=max_new_tokens,
temperature=temperature, top_k=top_k)
new = out[0, len(ids):].tolist()
if eot in new:
new = new[:new.index(eot)]
return tok.decode(new).strip()
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--prompt", required=True)
ap.add_argument("--temperature", type=float, default=0.5)
ap.add_argument("--top_k", type=int, default=20)
ap.add_argument("--max_new_tokens", type=int, default=70)
args = ap.parse_args()
model, tok = load()
print(rewrite(model, tok, args.prompt, args.temperature, args.top_k, args.max_new_tokens))
if __name__ == "__main__":
main()
|