"""第 24 课第 2 步：把输入交给已经训练好的模型。"""
import argparse
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(ROOT))
from poetry_gpt.common import DEFAULT_RUN
from poetry_gpt.generate import Writer, parse_keywords

parser = argparse.ArgumentParser(description="我的本地写诗命令")
parser.add_argument("--start", default="")
parser.add_argument("--keywords", default="")
parser.add_argument("--form", choices=["five", "seven"], default="five")
parser.add_argument("--device", choices=["auto", "cpu", "mps", "cuda"], default="auto")
args = parser.parse_args()

writer = Writer(DEFAULT_RUN / "best.pt", device=args.device)
poems = writer.write(
    start=args.start,
    keywords=parse_keywords(args.keywords),
    form=5 if args.form == "five" else 7,
)
print("\n".join(poems[0]["lines"]))