Skip to content
NLP

Word2Vec

Train and use word embeddings with gensim.

#word2vec#embeddings#gensim

Code

nlp
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)