Accept multiple values

This commit is contained in:
chris
2026-02-09 12:33:33 -05:00
parent 695f9dd687
commit 40207b7623

View File

@@ -12,8 +12,8 @@ def cli():
@cli.command() @cli.command()
@click.option("--length", default=50, help="Number of words to generate.") @click.option("--length", default=50, help="Number of words to generate.")
@click.option("--n", default=2, help="Number of words in the context window.") @click.option("--n", default=2, help="Number of words in the context window.")
@click.option("--text", type=click.Path(exists=True), help="Text file to use as training corpus.") @click.option("--text", type=click.Path(exists=True), multiple=True, help="Text file(s) to use as training corpus. Can be specified multiple times.")
@click.option("--gutenberg", help="NLTK Gutenberg corpus key (use --list-gutenberg to see available keys).") @click.option("--gutenberg", multiple=True, help="NLTK Gutenberg corpus key(s). Can be specified multiple times.")
@click.option("--list-gutenberg", is_flag=True, help="List available Gutenberg corpus keys.") @click.option("--list-gutenberg", is_flag=True, help="List available Gutenberg corpus keys.")
@click.option("--mbox", type=click.Path(exists=True), help="Mbox file to use for training.") @click.option("--mbox", type=click.Path(exists=True), help="Mbox file to use for training.")
@click.option("--prompt", help="Prompt to start generation.") @click.option("--prompt", help="Prompt to start generation.")
@@ -35,12 +35,16 @@ def generate(length, n, text, gutenberg, list_gutenberg, mbox, prompt, interact)
corpus = None corpus = None
if text: if text:
with open(text, "r") as f: corpus = []
corpus = f.read().split() for filepath in text:
with open(filepath, "r") as f:
corpus.extend(f.read().split())
elif gutenberg: elif gutenberg:
nltk.download("gutenberg", quiet=True) nltk.download("gutenberg", quiet=True)
from nltk.corpus import gutenberg as gutenberg_corpus from nltk.corpus import gutenberg as gutenberg_corpus
corpus = gutenberg_corpus.words(gutenberg) corpus = []
for key in gutenberg:
corpus.extend(gutenberg_corpus.words(key))
elif mbox: elif mbox:
mail_text = read_mail_text(mbox) mail_text = read_mail_text(mbox)
corpus = mail_text.split() corpus = mail_text.split()