BigramCollocationFinder
Pandas
DataFrame
Python
Natural Language Processing

How to apply a function BigramCollocationFinder to Pandas DataFrame

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

Pandas is excellent for tabular data, while NLTK's BigramCollocationFinder is designed for token sequences. Applying the two together means you first decide whether you want collocations per row, across the whole dataset, or within grouped subsets such as one document category at a time.

That distinction matters because BigramCollocationFinder does not operate directly on a DataFrame. It operates on a list of tokens, so the main task is preparing the text column correctly.

Applying the Finder to Each Row

If each row contains an independent document, review, or message, the simplest pattern is to tokenize each row and run the finder row by row.

python
1import pandas as pd
2from nltk.collocations import BigramAssocMeasures, BigramCollocationFinder
3from nltk.tokenize import wordpunct_tokenize
4
5df = pd.DataFrame(
6    {
7        "text": [
8            "machine learning models improve with clean data",
9            "clean data leads to better model quality",
10            "model quality depends on useful features",
11        ]
12    }
13)
14
15bigram_measures = BigramAssocMeasures()
16
17def top_bigrams(text):
18    tokens = [token.lower() for token in wordpunct_tokenize(text) if token.isalpha()]
19    finder = BigramCollocationFinder.from_words(tokens)
20    finder.apply_freq_filter(1)
21    return finder.nbest(bigram_measures.pmi, 3)
22
23df["top_bigrams"] = df["text"].apply(top_bigrams)
24print(df[["text", "top_bigrams"]])

This produces a list of top-scoring bigrams for each row. It is useful when you want row-level features or summaries.

Applying the Finder Across the Entire DataFrame

Sometimes you want collocations across the full corpus instead of inside each row. In that case, combine all row text into one token stream before creating the finder.

python
1import pandas as pd
2from nltk.collocations import BigramAssocMeasures, BigramCollocationFinder
3from nltk.tokenize import wordpunct_tokenize
4
5df = pd.DataFrame(
6    {
7        "text": [
8            "natural language processing needs tokenized text",
9            "language models learn from tokenized examples",
10            "clean text improves language analysis",
11        ]
12    }
13)
14
15tokens = []
16for text in df["text"].dropna():
17    row_tokens = [token.lower() for token in wordpunct_tokenize(text) if token.isalpha()]
18    tokens.extend(row_tokens)
19
20finder = BigramCollocationFinder.from_words(tokens)
21finder.apply_freq_filter(2)
22
23measures = BigramAssocMeasures()
24print(finder.nbest(measures.pmi, 10))

This version is better when you are building one vocabulary of interesting word pairs for the whole dataset.

Cleaning the Text Before Scoring

Collocation quality depends heavily on preprocessing. If you feed punctuation, stop words, and casing noise into the finder, the results often look trivial. A cleaner pipeline usually includes:

  • lowercasing
  • token filtering
  • stop word removal
  • minimum frequency thresholds

Here is a more realistic example:

python
1import pandas as pd
2from nltk.corpus import stopwords
3from nltk.collocations import BigramAssocMeasures, BigramCollocationFinder
4from nltk.tokenize import wordpunct_tokenize
5
6stop_words = set(stopwords.words("english"))
7
8def tokenize_for_collocations(text):
9    return [
10        token.lower()
11        for token in wordpunct_tokenize(text)
12        if token.isalpha() and token.lower() not in stop_words
13    ]
14
15df = pd.DataFrame(
16    {
17        "text": [
18            "The data science team builds data pipelines and data models",
19            "A science team needs reliable pipelines for analytics",
20        ]
21    }
22)
23
24tokens = []
25for text in df["text"]:
26    tokens.extend(tokenize_for_collocations(text))
27
28finder = BigramCollocationFinder.from_words(tokens)
29finder.apply_freq_filter(2)
30
31scored = finder.score_ngrams(BigramAssocMeasures().pmi)
32print(scored[:5])

score_ngrams returns both the bigram and its score, which is often more useful than only returning the top phrases.

Turning Bigrams into DataFrame Features

If your goal is machine learning, you may want a derived column rather than a printed list. One simple approach is to join the best bigrams into a readable feature column.

python
1def bigram_strings(text):
2    tokens = tokenize_for_collocations(text)
3    finder = BigramCollocationFinder.from_words(tokens)
4    return [" ".join(pair) for pair in finder.nbest(BigramAssocMeasures().pmi, 2)]
5
6df["bigram_feature"] = df["text"].apply(bigram_strings)
7print(df["bigram_feature"])

This can help during exploratory analysis. For production pipelines, you may instead convert collocations into structured features or counts.

Running the Finder Per Group

Sometimes the right unit is not each row and not the entire corpus, but a group such as product category, author, or label. In that case, aggregate text within each group and run the finder once per subset.

python
1grouped = df.groupby("category")["text"].apply(lambda values: " ".join(values))
2
3for category, text in grouped.items():
4    tokens = tokenize_for_collocations(text)
5    finder = BigramCollocationFinder.from_words(tokens)
6    print(category, finder.nbest(BigramAssocMeasures().pmi, 5))

This pattern is useful when you want category-specific phrase discovery instead of one global ranking.

Common Pitfalls

The most common mistake is passing an entire Pandas series directly to BigramCollocationFinder.from_words. The finder expects a flat token list, not a column of raw strings.

Another problem is forgetting that very short texts produce unstable scores. A row with three or four tokens does not contain enough context for meaningful collocation ranking.

Stop words can also dominate the results if you skip filtering. Bigram outputs such as "of the" or "in the" are statistically common but rarely useful.

Finally, remember that PMI tends to favor rare but exclusive pairs. If you want more common phrase-like results, compare different association measures instead of assuming one score fits every use case.

Summary

  • 'BigramCollocationFinder works on token sequences, not directly on a DataFrame.'
  • Use .apply(...) for per-row collocations or flatten the column for corpus-level analysis.
  • Clean tokenization and stop word filtering improve results substantially.
  • Frequency filters help remove noisy one-off pairs.
  • Choose the scoring method based on your goal, whether that is interpretability, ranking, or feature generation.

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.