wehub-resource-sync
66 lines
2.6 KiB
Python
66 lines
2.6 KiB
Python
"""
|
|
Chat / inference CLI for any stage checkpoint (base, SFT, DPO, PPO, GRPO).
|
|
|
|
Model dimensions are read from the checkpoint, so you only pass the path. Use the chat
|
|
template for instruction-tuned models, or --raw for base-model continuation.
|
|
|
|
One-shot:
|
|
PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/sft.pt --prompt "What is 13 + 29?"
|
|
PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/grpo.pt --prompt "..." --greedy
|
|
PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/base_pretrained.pt --raw --prompt "Once upon a time"
|
|
Interactive REPL (no --prompt):
|
|
PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/sft.pt
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
|
|
import torch
|
|
|
|
from src.post_training.inference import generate_reply, load_model_from_ckpt
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("--ckpt", required=True)
|
|
p.add_argument("--prompt", default=None, help="one-shot prompt; omit for interactive REPL")
|
|
p.add_argument("--system", default=None, help="optional system message (chat mode)")
|
|
p.add_argument("--raw", action="store_true", help="base-model continuation (no chat template)")
|
|
p.add_argument("--max_new_tokens", type=int, default=256)
|
|
p.add_argument("--temperature", type=float, default=0.8)
|
|
p.add_argument("--top_p", type=float, default=0.95)
|
|
p.add_argument("--top_k", type=int, default=None)
|
|
p.add_argument("--greedy", action="store_true", help="deterministic argmax decoding")
|
|
p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
|
|
args = p.parse_args()
|
|
|
|
model = load_model_from_ckpt(args.ckpt, args.device)
|
|
n = sum(p.numel() for p in model.parameters())
|
|
print(f"loaded {args.ckpt} ({n/1e6:.0f}M params) on {args.device} | "
|
|
f"mode={'raw' if args.raw else 'chat'} {'greedy' if args.greedy else f'T={args.temperature} top_p={args.top_p}'}")
|
|
|
|
def reply(text):
|
|
return generate_reply(model, text, device=args.device, system=args.system, raw=args.raw,
|
|
max_new_tokens=args.max_new_tokens, temperature=args.temperature,
|
|
top_p=args.top_p if args.top_p < 1 else None, top_k=args.top_k, greedy=args.greedy)
|
|
|
|
if args.prompt is not None:
|
|
print(reply(args.prompt))
|
|
return
|
|
|
|
print("Interactive chat (Ctrl-D / 'exit' to quit).")
|
|
while True:
|
|
try:
|
|
text = input("\nyou> ").strip()
|
|
except EOFError:
|
|
break
|
|
if text in ("exit", "quit"):
|
|
break
|
|
if text:
|
|
print("bot>", reply(text))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|