scikit-learn
CountVectorizer
Chi-Square
Natural Language Processing
Machine Learning

scikit learn Problems creating customized CountVectorizer and ChiSquare

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 and chi-square feature selection are a common pairing in text classification. The vectorizer turns text into token counts, and chi-square helps identify which terms are most associated with each class. Problems usually start when developers customize tokenization or preprocessing without realizing how that changes the feature matrix passed into feature selection.

What CountVectorizer Actually Produces

CountVectorizer converts a list of documents into a sparse matrix where each column represents a token and each row represents a document. By default it lowercases text, tokenizes using a regular expression, and counts occurrences.

That means customization can affect several layers at once:

  • how raw text is cleaned
  • how tokens are split
  • whether stop words are removed
  • whether unigrams, bigrams, or both are emitted

If your tokenizer returns inconsistent output, chi-square is not the real problem. The features themselves are unstable.

A Clean Pipeline With Custom Tokenization

The safest pattern is to build the vectorizer and selector inside a single scikit-learn pipeline. That keeps training and prediction behavior aligned.

python
1import re
2from sklearn.feature_extraction.text import CountVectorizer
3from sklearn.feature_selection import SelectKBest, chi2
4from sklearn.linear_model import LogisticRegression
5from sklearn.pipeline import Pipeline
6
7
8def tokenizer(text: str) -> list[str]:
9    return re.findall(r"[a-zA-Z']+", text.lower())
10
11
12X = [
13    "Fast delivery and great support",
14    "Support team fixed my problem quickly",
15    "Terrible delay and broken packaging",
16    "Broken item and slow delivery",
17]
18
19y = [1, 1, 0, 0]
20
21pipeline = Pipeline([
22    ("vectorizer", CountVectorizer(tokenizer=tokenizer, ngram_range=(1, 2))),
23    ("selector", SelectKBest(score_func=chi2, k=6)),
24    ("model", LogisticRegression(max_iter=1000)),
25])
26
27pipeline.fit(X, y)
28prediction = pipeline.predict(["great support and fast fix"])
29print(prediction[0])

This example is intentionally small, but the structure is the important part. Tokenization, feature selection, and modeling all run in one consistent path.

Why Chi-Square Fails in Custom Setups

Chi-square measures dependence between a non-negative feature and a target label. For text classification, token counts fit that requirement naturally. But custom preprocessing often breaks one of the assumptions.

Common failure modes include:

  • passing transformed values that are negative
  • fitting the vectorizer on one dataset and the selector on a differently processed dataset
  • using a tokenizer that returns unexpected objects instead of strings
  • forgetting that chi-square should be fit only on training data, not on the full dataset

If you use a custom transformer before chi-square, ensure the output remains non-negative. Raw counts and term frequencies are fine. Centered or standardized values are not.

Inspecting Features After Vectorization

When debugging, inspect the learned vocabulary and matrix shape before blaming feature selection.

python
1from sklearn.feature_extraction.text import CountVectorizer
2
3
4def tokenizer(text: str) -> list[str]:
5    return text.lower().split()
6
7
8vectorizer = CountVectorizer(tokenizer=tokenizer)
9X = vectorizer.fit_transform([
10    "red apple red",
11    "green apple",
12    "green banana",
13])
14
15print(vectorizer.get_feature_names_out())
16print(X.toarray())

That output tells you exactly which tokens survived preprocessing and how often they appear. If a domain-specific token is missing, the issue is usually with tokenization, lowercasing, accent stripping, stop-word removal, or the regular expression used by the analyzer.

Choosing Between a Tokenizer and an Analyzer

Many customizations that developers implement as a custom tokenizer belong in preprocessor, token_pattern, or ngram_range instead. The more of scikit-learn's built-in pipeline you keep, the less surprising the behavior becomes.

Use a custom tokenizer when token boundaries are truly domain-specific, such as chemical tokens, source code symbols, or special identifiers. Otherwise, prefer simpler configuration:

python
1vectorizer = CountVectorizer(
2    lowercase=True,
3    stop_words="english",
4    ngram_range=(1, 2),
5    token_pattern=r"(?u)\b\w\w+\b",
6)

That version is easier to serialize, easier to reason about, and usually easier to reproduce across training jobs.

Interpreting Chi-Square Scores

Chi-square does not tell you whether a token is globally useful in every sense. It tells you how strongly the feature distribution differs across labels. A token can have a high chi-square score because it is very common in one class and rare in another.

After fitting SelectKBest, inspect the selected features:

python
1vectorizer = CountVectorizer()
2X_counts = vectorizer.fit_transform(X)
3selector = SelectKBest(chi2, k=3)
4selector.fit(X_counts, y)
5
6selected = vectorizer.get_feature_names_out()[selector.get_support()]
7print(selected)

This is often the fastest way to validate that your custom vectorization is producing meaningful signals instead of artifacts.

Common Pitfalls

A frequent mistake is overriding tokenizer while forgetting that token_pattern may no longer matter the way you expect. Read the vectorizer arguments as a system, not as independent switches.

Another bug is doing manual preprocessing before the pipeline and then slightly different preprocessing during inference. That creates train-predict skew.

Do not run chi-square on data that contains negative values. It is built for count-like or frequency-like non-negative features.

Also avoid fitting the vectorizer, selector, and model separately unless you have a strong reason. Pipelines prevent subtle ordering mistakes and make cross-validation much safer.

Summary

  • 'CountVectorizer creates the sparse token-count matrix that chi-square evaluates.'
  • Most customization bugs come from unstable preprocessing, not from chi-square itself.
  • Keep vectorization, selection, and modeling in a single scikit-learn pipeline.
  • Chi-square expects non-negative feature values.
  • Inspect vocabulary and matrix output when debugging custom tokenization.
  • Prefer built-in vectorizer options over a custom tokenizer unless the domain really requires it.

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.