| |
| """ |
| generate.py — CLI for the inference.py library bundled in this repo. |
| |
| One-shot: |
| python generate.py --prompt "<s><|yor_Latn|>Kaabo..." -n 4 |
| |
| Interactive REPL: |
| python generate.py --interactive |
| |
| Prompt format: prefix with <s> (BOS) and a language token like <|yor_Latn|>. |
| Available language tokens are listed in tokenizer.json. |
| """ |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| from inference import load_model, generate |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description=__doc__, |
| formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--repo-dir", default=".", help="Path to the repo (default: cwd).") |
| ap.add_argument("--prompt", default=None, help="Single prompt to complete.") |
| ap.add_argument("--interactive", action="store_true", |
| help="REPL mode. Type prompts; empty line or Ctrl-D exits.") |
| ap.add_argument("-n", "--num-sequences", type=int, default=1) |
| ap.add_argument("--max-new-tokens", type=int, default=128) |
| ap.add_argument("--temperature", type=float, default=0.8) |
| ap.add_argument("--top-p", type=float, default=0.9) |
| ap.add_argument("--seed", type=int, default=42) |
| ap.add_argument("--device", default=None, |
| help="cuda / cpu (default: cuda if available).") |
| args = ap.parse_args() |
|
|
| if not args.prompt and not args.interactive: |
| print("[error] specify --prompt or --interactive", file=sys.stderr); sys.exit(1) |
|
|
| print(f"[load] {args.repo_dir}") |
| model, tokenizer, config = load_model(args.repo_dir, device=args.device) |
| n_params = sum(p.numel() for p in model.parameters()) |
| print(f"[load] {n_params/1e6:.1f}M params, dtype={next(model.parameters()).dtype}, " |
| f"vocab={config.vocab_size}, T={config.max_seq_len}") |
|
|
| def run_one(prompt: str): |
| out = generate( |
| model, tokenizer, prompt, |
| max_new_tokens=args.max_new_tokens, |
| temperature=args.temperature, top_p=args.top_p, |
| num_sequences=args.num_sequences, seed=args.seed, |
| ) |
| for i, c in enumerate(out): |
| print(f"\n[sample {i}]") |
| print(c) |
|
|
| if args.prompt: |
| run_one(args.prompt) |
| return |
|
|
| print("\nInteractive mode. Empty line or Ctrl-D to exit.") |
| print('Prompt format: <s><|lang_Script|>your text here') |
| while True: |
| try: |
| p = input("\n>>> ").rstrip() |
| except (EOFError, KeyboardInterrupt): |
| print(); break |
| if not p: |
| break |
| run_one(p) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
