word2vec.py

"""
Loads a textfile, builds a Word2Vec model, and prints similarity of words.
"""
import urllib.request
import nltk
from nltk.tokenize import sent_tokenize, word_tokenize
from gensim.models import Word2Vec

# Source: Project Gutenberg's Alice's Adventures in Wonderland.
CORPUS_TEXT_URL = 'https://www.gutenberg.org/files/11/11.txt'


def main():
    setup()

    # Gets training data from corpus.
    data = build_data(get_corpus())

    # Trains the word2vec model.
    model = Word2Vec(data, min_count=1, size=25, window=5, sg=1)

    # Gets the most similar words with 'alice'.
    print(model.similarity('alice', 'king'))
    print(model.similarity('alice', 'tree'))


def setup():
    nltk.download('punkt')


def get_corpus():
    return ' '.join(
        urllib.request.urlopen(CORPUS_TEXT_URL).read().decode("utf-8")
        .replace('\n', ' ').split()
    )


def build_data(corpus):
    return [
        [word.lower() for word in word_tokenize(sentence)]
        for sentence in sent_tokenize(corpus)
    ]


if __name__ == '__main__':
    main()