Word2Vec

Code

from gensim.models import Word2Vec, KeyedVectors

sentences = [
    ["the", "cat", "sat", "on", "the", "mat"],
    ["the", "dog", "sat", "on", "the", "log"],
    ["cats", "and", "dogs", "are", "pets"],
    ["i", "love", "my", "cat"],
]

model = Word2Vec(sentences, vector_size=100, window=5, min_count=1,
                 workers=4, sg=1, epochs=50)

# Vector and similarity
vec = model.wv["cat"]
print(vec.shape)

sim = model.wv.most_similar("cat", topn=3)
print(sim)

# Analogy: king - man + woman ~= queen
# analogy = model.wv.most_similar(positive=["king", "woman"], negative=["man"])

# Save and load
model.wv.save_word2vec_format("vecs.txt", binary=False)
w2v = KeyedVectors.load_word2vec_format("vecs.txt", binary=False)

# Pretrained vectors
# pre = KeyedVectors.load_word2vec_format("GoogleNews-vectors-negative300.bin", binary=True)

Explanation

Word2Vec learns dense vector embeddings by predicting context words (or vice versa) over a corpus. vector_size sets dimensionality, sg=1 uses skip-gram and sg=0 uses CBOW, and window controls context radius. The resulting KeyedVectors support similarity queries and analogy via vector arithmetic.

More NLP Snippets