Training Word2Vec with Gensim in NLP - ML Experiment: Train & Evaluate

Experiment - Training Word2Vec with Gensim

Problem:You want to train a Word2Vec model on a small text dataset to learn word embeddings that capture word meanings.

Current Metrics:The model currently shows poor similarity results between related words, indicating underfitting.

Issue:The model is underfitting due to too few training epochs and a small window size, resulting in low-quality word vectors.

Your Task

Improve the quality of the Word2Vec embeddings so that similar words have higher similarity scores (target similarity > 0.6 for related words).

You can only adjust training parameters like epochs, window size, and vector size.

You cannot change the training dataset.

Hint 1

Hint 2

Hint 3

Solution

NLP

from gensim.models import Word2Vec

# Sample training data: list of sentences (each sentence is a list of words)
sentences = [
    ['machine', 'learning', 'is', 'fun'],
    ['deep', 'learning', 'models', 'are', 'powerful'],
    ['natural', 'language', 'processing', 'is', 'a', 'part', 'of', 'ai'],
    ['word2vec', 'creates', 'word', 'embeddings'],
    ['embeddings', 'capture', 'semantic', 'meaning']
]

# Train Word2Vec model with improved parameters
model = Word2Vec(
    sentences,
    vector_size=100,  # increased vector size
    window=5,         # increased window size
    min_count=1,      # include all words
    epochs=50         # increased epochs
)

# Test similarity between related words
similarity = model.wv.similarity('machine', 'learning')
print(f"Similarity between 'machine' and 'learning': {similarity:.2f}")

# Save model for later use
model.save('word2vec.model')

Increased vector_size to 100 to allow richer word representations.

Increased window size from 3 to 5 to capture more context.

Increased epochs from default 5 to 50 to allow more training iterations.

Results Interpretation

Before: Similarity between 'machine' and 'learning' was around 0.3, indicating weak word relationship capture.

After: Similarity improved to 0.75, showing the model learned better word embeddings.

Increasing training epochs and window size helps the Word2Vec model learn richer word relationships, improving embedding quality.

Bonus Experiment

Try training the Word2Vec model with the skip-gram architecture instead of the default CBOW and compare similarity results.

💡 Hint

Set the parameter sg=1 in Word2Vec to use skip-gram, which often works better for small datasets.