| """Text generation CLI. |
| |
| python -m src.generate --prompt "ROMEO:" --max-new-tokens 200 --temperature 0.8 |
| python -m src.generate --fixed-prompts # runs the 5 fixed prompts from the config |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import torch |
|
|
| from src.utils import get_device, load_checkpoint, load_config, resolve |
|
|
|
|
| def generate_text(model, tok, prompt: str, max_new_tokens: int = 200, temperature: float = 0.8, |
| top_k=None, top_p=None, greedy: bool = False, device="cpu") -> str: |
| ids = torch.tensor([tok.encode(prompt)], dtype=torch.long, device=device) |
| out = model.generate(ids, max_new_tokens, temperature=temperature, top_k=top_k, top_p=top_p, greedy=greedy) |
| return tok.decode(out[0].tolist()) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description="Generate text with a trained MiniGPT") |
| ap.add_argument("--checkpoint", default="checkpoints/best.pt") |
| ap.add_argument("--prompt", default="ROMEO:") |
| ap.add_argument("--max-new-tokens", type=int, default=200) |
| ap.add_argument("--temperature", type=float, default=0.8) |
| ap.add_argument("--top-k", type=int, default=None) |
| ap.add_argument("--top-p", type=float, default=None) |
| ap.add_argument("--greedy", action="store_true") |
| ap.add_argument("--seed", type=int, default=1337) |
| ap.add_argument("--fixed-prompts", action="store_true", help="run the fixed prompts from the config") |
| ap.add_argument("--device", default="auto") |
| args = ap.parse_args() |
|
|
| device = get_device(args.device) |
| torch.manual_seed(args.seed) |
| model, tok, ckpt = load_checkpoint(args.checkpoint, device) |
| n = model.count_parameters() |
| prompts = load_config()["generate"]["fixed_prompts"] if args.fixed_prompts else [args.prompt] |
| blocks = [] |
| for p in prompts: |
| text = generate_text(model, tok, p, args.max_new_tokens, args.temperature, args.top_k, args.top_p, args.greedy, device) |
| block = (f"Prompt: {p!r}\nGenerated: {text}\n" |
| f"model parameters: {n / 1e6:.2f}M\ncontext length: {model.context_length}\n" |
| f"temperature: {'greedy' if args.greedy else args.temperature}\n" + "-" * 60) |
| print(block) |
| blocks.append(block) |
| if args.fixed_prompts: |
| out = resolve("outputs/samples/fixed_prompt_samples.txt") |
| out.parent.mkdir(parents=True, exist_ok=True) |
| out.write_text("\n".join(blocks), encoding="utf-8") |
| print(f"saved to {out}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
