Accept multiple values
This commit is contained in:
14
tlm/cli.py
14
tlm/cli.py
@@ -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()
|
||||||
|
|||||||
Reference in New Issue
Block a user