llm/generate.py

配套源码，运行方法见同目录 README。返回实验总览。

"""Generate with the saved CPU Transformer; no training occurs here."""
import argparse
from pathlib import Path
from model import TinyTransformer, generate


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", type=Path, default=Path(__file__).with_name("runs") / "model.npz")
    parser.add_argument("--prompt", default="the fox ")
    parser.add_argument("--tokens", type=int, default=120)
    parser.add_argument("--temperature", type=float, default=.8)
    parser.add_argument("--top-k", type=int, default=10)
    parser.add_argument("--seed", type=int, default=9)
    args = parser.parse_args()
    model, tokenizer = TinyTransformer.load(args.checkpoint)
    try:
        prompt = tokenizer.encode(args.prompt)
    except KeyError as error:
        parser.error(f"character {error} is outside the training alphabet; use --tokenizer bpe when training for arbitrary UTF-8 input")
    ids = generate(model, prompt, args.tokens, args.temperature, args.top_k, args.seed)
    print(tokenizer.decode(ids))


if __name__ == "__main__":
    main()