| """ |
| Try a model by hand. |
| |
| python chat.py --ckpt checkpoints/dpo.pt --words dragon,cookie,rain |
| python chat.py --ckpt checkpoints/base.pt --text "Once upon a time" |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import os |
|
|
| import torch |
|
|
| from common import build_prompt, get_device |
| from common import load_model |
| from config import get_config |
| from tokenizer import BPETokenizer |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ckpt", default="checkpoints/dpo.pt") |
| ap.add_argument("--words", help="required words, comma-separated (instruction mode)") |
| ap.add_argument("--text", help="prompt text (continuation mode, suited to the base model)") |
| ap.add_argument("--n", type=int, default=2) |
| ap.add_argument("--temperature", type=float, default=0.7) |
| ap.add_argument("--max_tokens", type=int, default=200) |
| args = ap.parse_args() |
|
|
| cfg = get_config() |
| device = get_device() |
| tok = BPETokenizer.load(os.path.join(cfg.data.data_dir, "tokenizer.json")) |
| model, mcfg = load_model(args.ckpt, device) |
| model.eval() |
|
|
| if args.words: |
| prompt = build_prompt([w.strip() for w in args.words.split(",") if w.strip()]) |
| else: |
| prompt = args.text or "Once upon a time" |
| ids = tok.encode(prompt, allow_special=bool(args.words)) |
| x = torch.tensor([ids] * args.n, device=device) |
| budget = min(args.max_tokens, mcfg.context_len - len(ids)) |
| out, _ = model.generate(x, budget, temperature=args.temperature, top_k=50, eos_id=tok.eos_id, valid_vocab=tok.vocab_size) |
| for i in range(args.n): |
| text = tok.decode([t for t in out[i].tolist() if t != tok.eos_id]) |
| print(f"\n--- sample {i + 1} ---\n{text.strip()}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
