CountVectorizer
machine learning
feature weighting
text analysis
natural language processing

How to increase weight of a word for CountVectorizer

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

CountVectorizer from scikit-learn converts text into a matrix of token counts where each word gets equal weight. To increase the importance of specific words, you can multiply their columns in the count matrix by a weight factor, use TfidfVectorizer (which automatically downweights common words and upweights rare ones), create a custom vocabulary with repeated tokens, or build a custom transformer. The right approach depends on whether you want domain-specific boosting or automatic importance weighting.

Basic CountVectorizer

python
1from sklearn.feature_extraction.text import CountVectorizer
2
3corpus = [
4    "machine learning is great",
5    "deep learning uses neural networks",
6    "machine learning and deep learning are related",
7]
8
9vectorizer = CountVectorizer()
10X = vectorizer.fit_transform(corpus)
11
12print(vectorizer.get_feature_names_out())
13# ['and', 'are', 'deep', 'great', 'is', 'learning', 'machine', 'networks', 'neural', 'related', 'uses']
14
15print(X.toarray())
16# [[0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0],
17#  [0, 0, 1, 0, 0, 1, 0, 1, 1, 0, 1],
18#  [1, 1, 1, 0, 0, 2, 1, 0, 0, 1, 0]]

Every word gets a count of how many times it appears. "learning" appearing twice in the third document gets a count of 2, but all words are treated equally in terms of importance.

Method 1: Multiply Specific Columns by a Weight

python
1import numpy as np
2from sklearn.feature_extraction.text import CountVectorizer
3from scipy.sparse import csr_matrix
4
5corpus = [
6    "python is great for machine learning",
7    "java is used in enterprise systems",
8    "python and java are programming languages",
9]
10
11vectorizer = CountVectorizer()
12X = vectorizer.fit_transform(corpus)
13
14# Define weights for specific words
15word_weights = {
16    "python": 3.0,   # Boost Python mentions
17    "machine": 2.0,  # Boost machine learning terms
18    "learning": 2.0,
19}
20
21# Apply weights to the count matrix
22feature_names = vectorizer.get_feature_names_out()
23weight_array = np.ones(len(feature_names))
24
25for word, weight in word_weights.items():
26    if word in feature_names:
27        idx = list(feature_names).index(word)
28        weight_array[idx] = weight
29
30# Multiply columns by weights
31X_weighted = X.multiply(weight_array)
32print(X_weighted.toarray())

This approach directly scales specific feature columns. It is the most straightforward way to boost known important terms.

Method 2: Use TfidfVectorizer Instead

python
1from sklearn.feature_extraction.text import TfidfVectorizer
2
3corpus = [
4    "the cat sat on the mat",
5    "the dog sat on the log",
6    "the cat and the dog are friends",
7]
8
9# TF-IDF automatically downweights common words ("the", "on")
10# and upweights distinctive words ("mat", "log", "friends")
11tfidf = TfidfVectorizer()
12X = tfidf.fit_transform(corpus)
13
14# Show TF-IDF weights
15feature_names = tfidf.get_feature_names_out()
16for doc_idx in range(len(corpus)):
17    print(f"\nDocument {doc_idx}:")
18    for word_idx in X[doc_idx].nonzero()[1]:
19        print(f"  {feature_names[word_idx]}: {X[doc_idx, word_idx]:.3f}")

TfidfVectorizer computes Term Frequency-Inverse Document Frequency. Words that appear in many documents (like "the") get low weight, while words unique to few documents get high weight. This is the standard approach for automatic importance weighting.

Method 3: Custom Transformer with Weights

python
1from sklearn.base import BaseEstimator, TransformerMixin
2from sklearn.feature_extraction.text import CountVectorizer
3import numpy as np
4
5class WeightedCountVectorizer(BaseEstimator, TransformerMixin):
6    def __init__(self, word_weights=None, **cv_params):
7        self.word_weights = word_weights or {}
8        self.cv_params = cv_params
9        self.vectorizer = CountVectorizer(**cv_params)
10
11    def fit(self, X, y=None):
12        self.vectorizer.fit(X)
13        feature_names = self.vectorizer.get_feature_names_out()
14        self.weights_ = np.ones(len(feature_names))
15        for word, weight in self.word_weights.items():
16            if word in feature_names:
17                idx = list(feature_names).index(word)
18                self.weights_[idx] = weight
19        return self
20
21    def transform(self, X):
22        counts = self.vectorizer.transform(X)
23        return counts.multiply(self.weights_)
24
25    def get_feature_names_out(self):
26        return self.vectorizer.get_feature_names_out()
27
28# Usage
29wv = WeightedCountVectorizer(word_weights={"python": 3, "learning": 2})
30X = wv.fit_transform(corpus)

A custom transformer integrates into scikit-learn pipelines and applies weights consistently during both training and prediction.

Method 4: Using Pipeline with Custom Weights

python
1from sklearn.pipeline import Pipeline
2from sklearn.preprocessing import FunctionTransformer
3from sklearn.linear_model import LogisticRegression
4
5def apply_weights(X):
6    """Apply domain-specific weights to feature matrix."""
7    weights = np.ones(X.shape[1])
8    # Boost columns 0 and 3 (known important features)
9    weights[0] = 2.0
10    weights[3] = 3.0
11    return X.multiply(weights)
12
13pipeline = Pipeline([
14    ('vectorizer', CountVectorizer()),
15    ('weighter', FunctionTransformer(apply_weights, accept_sparse=True)),
16    ('classifier', LogisticRegression()),
17])
18
19pipeline.fit(train_texts, train_labels)
20predictions = pipeline.predict(test_texts)

Method 5: Vocabulary Duplication Trick

python
1from sklearn.feature_extraction.text import CountVectorizer
2
3def boost_words(text, boost_words, multiplier=3):
4    """Repeat specific words to increase their count."""
5    words = text.split()
6    boosted = []
7    for word in words:
8        if word.lower() in boost_words:
9            boosted.extend([word] * multiplier)
10        else:
11            boosted.append(word)
12    return ' '.join(boosted)
13
14corpus = [
15    "python is great for data science",
16    "java is used for enterprise apps",
17]
18
19# Boost "python" and "data" by repeating them
20boosted_corpus = [boost_words(doc, {"python", "data"}, multiplier=3) for doc in corpus]
21# "python python python is great for data data data science"
22
23vectorizer = CountVectorizer()
24X = vectorizer.fit_transform(boosted_corpus)

Repeating words in the input text increases their count in the feature matrix. This is a hack — prefer the column-weighting approach for production code.

Common Pitfalls

  • Applying weights after fitting the model: Weights must be applied during both training and prediction. If you weight the training data but not the test data, the model sees different feature distributions and performs poorly.
  • Choosing arbitrary weight values: Weight factors should be validated via cross-validation. Boosting a word by 10x may cause overfitting. Start with small multipliers (2-3x) and tune on validation data.
  • Using CountVectorizer when TF-IDF suffices: If your goal is to reduce the impact of common words and boost rare ones, TfidfVectorizer does this automatically without manual weight tuning.
  • Vocabulary duplication changing model semantics: Repeating words in the text changes the total token count per document, which affects normalization. The column-weighting approach preserves the original document structure.
  • Not accounting for word weights in feature importance: If you weight columns, feature importance scores from the model reflect the weighted values, not the raw counts. Document the weights applied so model interpretation remains valid.

Summary

  • Multiply specific columns in the count matrix by a weight factor for direct control
  • Use TfidfVectorizer for automatic importance weighting based on document frequency
  • Build a custom WeightedCountVectorizer transformer for scikit-learn pipeline integration
  • Apply weights consistently during both training and prediction to avoid data leakage
  • Validate weight choices with cross-validation rather than using arbitrary multipliers
  • Prefer column weighting over vocabulary duplication for cleaner, more maintainable code

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.