62 lines
1.6 KiB
Python
62 lines
1.6 KiB
Python
import click
|
|
from . import wordplay
|
|
|
|
SIZES = {
|
|
'small': 'GloVe Wikipedia 100-dim (~128 MB)',
|
|
'medium': 'GloVe Wikipedia 300-dim (~376 MB)',
|
|
'large': 'Word2Vec Google News 300-dim (~1.6 GB)',
|
|
}
|
|
|
|
|
|
@click.group()
|
|
@click.option(
|
|
'--size',
|
|
type=click.Choice(list(SIZES)),
|
|
default='small',
|
|
show_default=True,
|
|
help='Embedding model to use:\n\n' + '\n'.join(f' {k}: {v}' for k, v in SIZES.items()),
|
|
)
|
|
@click.pass_context
|
|
def cli(ctx, size):
|
|
"""Wordplay - word embedding experiments."""
|
|
ctx.ensure_object(dict)
|
|
ctx.obj['model'] = wordplay.load_embeddings(size)
|
|
|
|
|
|
@cli.command()
|
|
@click.argument('word')
|
|
@click.option('-n', default=5, show_default=True, help='Number of synonyms to return.')
|
|
@click.pass_obj
|
|
def synonyms(obj, word, n):
|
|
"""Find the n most similar words to WORD."""
|
|
click.echo(wordplay.synonyms(obj['model'], word, n))
|
|
|
|
|
|
@cli.command()
|
|
@click.argument('word1')
|
|
@click.argument('word2')
|
|
@click.pass_obj
|
|
def average(obj, word1, word2):
|
|
"""Return a word which is the average of the two given words."""
|
|
click.echo(wordplay.average(obj['model'], word1, word2))
|
|
|
|
|
|
@cli.command()
|
|
@click.argument('words', nargs=-1, required=True)
|
|
@click.pass_obj
|
|
def outlier(obj, words):
|
|
"""Given a list of words, identifies the word which is the outlier."""
|
|
click.echo(wordplay.outlier(obj['model'], words))
|
|
|
|
|
|
@cli.command()
|
|
@click.argument('words', nargs=-1, required=True)
|
|
@click.pass_obj
|
|
def sort(obj, words):
|
|
"""Given a list of words, sorts them in order."""
|
|
click.echo(wordplay.sort(obj['model'], words))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
cli()
|