Upload generate.py with huggingface_hub

generate.py

@@ -0,0 +1,53 @@
+"""
+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()