word2vec
word classification
machine learning
natural language processing
NLP

Using word2vec to classify words in categories

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

Word2Vec does not classify words by itself. It learns vector representations where semantically similar words end up close together, and you then use those vectors as features for a separate classification rule or model. Once that distinction is clear, Word2Vec becomes a practical building block for category assignment.

What Word2Vec Gives You

Word2Vec converts tokens into dense numeric vectors. After training, words with similar contexts often have nearby embeddings. That means you can compare words with cosine similarity or feed their vectors into a classifier.

In practice, there are two common approaches:

  • nearest prototype or centroid for each category
  • supervised classifier trained on labeled word vectors

If you only have a few labeled examples per category, centroid classification is often enough.

Train a Small Word2Vec Model

This example uses gensim to learn embeddings from a toy corpus. The data is intentionally small so the code is readable, but in real work you usually need a much larger corpus.

python
1from gensim.models import Word2Vec
2
3sentences = [
4    ["dog", "puppy", "pet", "animal"],
5    ["cat", "kitten", "pet", "animal"],
6    ["apple", "banana", "fruit", "food"],
7    ["orange", "grape", "fruit", "food"],
8    ["car", "truck", "vehicle", "road"],
9    ["bus", "train", "vehicle", "transport"],
10]
11
12model = Word2Vec(
13    sentences=sentences,
14    vector_size=20,
15    window=2,
16    min_count=1,
17    workers=1,
18    epochs=200,
19)
20
21print(model.wv.similar_by_word("dog", topn=3))

The important artifact is model.wv[word], which returns the embedding vector for a token.

Classify with Category Centroids

A simple strategy is to average the vectors of example words for each category and assign a new word to the nearest centroid.

python
1import numpy as np
2
3categories = {
4    "animals": ["dog", "cat", "puppy", "kitten"],
5    "fruits": ["apple", "banana", "orange", "grape"],
6    "vehicles": ["car", "truck", "bus", "train"],
7}
8
9
10def centroid(words):
11    return np.mean([model.wv[w] for w in words], axis=0)
12
13
14centroids = {name: centroid(words) for name, words in categories.items()}
15
16
17def classify_word(word):
18    vec = model.wv[word]
19    best_label = None
20    best_score = -1.0
21
22    for label, center in centroids.items():
23        score = np.dot(vec, center) / (np.linalg.norm(vec) * np.linalg.norm(center))
24        if score > best_score:
25            best_score = score
26            best_label = label
27    return best_label, best_score
28
29
30print(classify_word("kitten"))
31print(classify_word("truck"))

This is useful when categories are conceptually coherent and you do not have much labeled data.

Use a Supervised Classifier When Labels Matter

If you have labeled words, you can train a standard classifier on top of Word2Vec features. That works better when category boundaries are subtle or when one centroid per category is too simplistic.

python
1from sklearn.linear_model import LogisticRegression
2
3labeled_words = ["dog", "cat", "apple", "banana", "car", "truck"]
4labels = ["animal", "animal", "fruit", "fruit", "vehicle", "vehicle"]
5X = [model.wv[w] for w in labeled_words]
6
7clf = LogisticRegression(max_iter=1000)
8clf.fit(X, labels)
9
10print(clf.predict([model.wv["bus"]])[0])

The advantage here is that the classifier can learn more flexible decision boundaries than pure nearest-centroid matching.

Handle Out-of-Vocabulary Words

Word2Vec has an important limitation: if a word was not seen during training, there is no vector for it. That means classification fails for out-of-vocabulary items unless you add a fallback strategy.

Common options include:

  • retrain or continue training with a larger corpus
  • normalize tokens aggressively before lookup
  • switch to subword models such as FastText when rare words matter

For production systems, this limitation often determines whether Word2Vec is the right tool at all.

Common Pitfalls

  • Expecting Word2Vec alone to produce category labels without any downstream rule or classifier.
  • Training on a tiny or unrepresentative corpus and expecting stable semantic structure.
  • Ignoring out-of-vocabulary words until inference starts failing.
  • Treating similar context as identical meaning, even when categories are task-specific.
  • Evaluating the classifier only on obvious examples and missing edge words with multiple senses.

Summary

  • Word2Vec learns embeddings, not labels.
  • To classify words, use embeddings with a centroid rule or a supervised classifier.
  • 'gensim makes it straightforward to train a small model and access vectors.'
  • Out-of-vocabulary handling is a real design constraint for Word2Vec systems.
  • Good category accuracy depends as much on corpus quality and labeling strategy as on the embedding model itself.

Course illustration
Course illustration

All Rights Reserved.