natural language toolkit collocations and association measures c 20012023 nltk project joel nothman jnothmanstudent usyd edu au url https www nltk org for license information see license txt tools to identify collocations words that often appear consecutively within corpora they may also be used to find other associations between word occurrences see manning and schutze ch 5 at https nlp stanford edufsnlppromocolloc pdf and the text nsp perl package at http ngram sourceforge net finding collocations requires first calculating the frequencies of words and their appearance in the context of other words often the collection of words will then requiring filtering to only retain useful content terms each ngram of words may then be scored according to some association measure in order to determine the relative likelihood of each ngram being a collocation the bigramcollocationfinder and trigramcollocationfinder classes provide these functionalities dependent on being provided a function which scores a ngram given appropriate frequency counts a number of standard association measures are provided in bigrammeasures and trigrammeasures possible todos consider the distinction between fx and fx and whether our approximation is good enough for fragmented data and mention it add a ngram collocation finder with measures which only utilise ngram and unigram counts rawfreq pmi studentt these two unused imports are referenced in collocations doctest an abstract base class for collocation finders whose purpose is to collect collocation candidate frequencies filter and rank them as a minimum collocation finders require the frequencies of each word in a corpus and the joint frequency of word tuples this data should be provided through nltk probability freqdist objects or an identical interface pad the document with the place holder according to the windowsize constructs a collocation finder given a collection of documents each of which is a list or iterable of tokens return cls fromwordsitertools chaindocuments generic filter removes ngrams from the frequency distribution if the function returns true when passed an ngram tuple removes candidate ngrams which have frequency less than minfreq self applyfilterlambda ng freq freq minfreq def applyngramfilterself fn self applyfilterlambda ng f fnng def applywordfilterself fn self applyfilterlambda ng f anyfnw for w in ng def scorengramsself scorefn for tup in self ngramfd score self scorengramscorefn tup if score is not none yield tup score def scorengramsself scorefn return sortedself scorengramsscorefn keylambda t t1 t0 def nbestself scorefn n returns a sequence of ngrams ordered by decreasing score whose scores each exceed the given minimum score a tool for the finding and ranking of bigram collocations or other association measures it is often useful to use fromwords rather than constructing an instance directly construct a bigramcollocationfinder given freqdists for appearances of words and possibly noncontiguous bigrams construct a bigramcollocationfinder for all bigrams in the given sequence when windowsize 2 count noncontiguous bigrams in the style of church and hanks s 1990 association ratio returns the score for a given bigram using the given scoring function following church and hanks 1990 counts are scaled by a factor of 1windowsize 1 a tool for the finding and ranking of trigram collocations or other association measures it is often useful to use fromwords rather than constructing an instance directly construct a trigramcollocationfinder given freqdists for appearances of words bigrams two words with any word between them and trigrams construct a trigramcollocationfinder for all trigrams in the given sequence constructs a bigram collocation finder with the bigram and unigram data from this finder note that this does not include any filtering applied to this finder returns the score for a given trigram using the given scoring function a tool for the finding and ranking of quadgram collocations or other association measures it is often useful to use fromwords rather than constructing an instance directly construct a quadgramcollocationfinder given freqdists for appearances of words bigrams trigrams two words with one word and two words between them three words with a word between them in both variations finds bigram collocations in the files of the webtext corpus from nltk metrics import bigramassocmeasures ranksfromscores spearmancorrelation if scorer is none scorer bigramassocmeasures likelihoodratio if comparescorer is none comparescorer bigramassocmeasures rawfreq from nltk corpus import stopwords webtext ignoredwords stopwords wordsenglish wordfilter lambda w lenw 3 or w lower in ignoredwords for file in webtext fileids words word lower for word in webtext wordsfile cf bigramcollocationfinder fromwordswords cf applyfreqfilter3 cf applywordfilterwordfilter corr spearmancorrelation ranksfromscorescf scorengramsscorer ranksfromscorescf scorengramscomparescorer printfile printt jointup for tup in cf nbestscorer 15 printft correlation to comparescorer name corr 0 4f slows down loading too much bigrammeasures bigramassocmeasures trigrammeasures trigramassocmeasures if name main import sys from nltk metrics import bigramassocmeasures try scorer evalbigramassocmeasures sys argv1 except indexerror scorer none try comparescorer evalbigramassocmeasures sys argv2 except indexerror comparescorer none demoscorer comparescorer all bigramcollocationfinder trigramcollocationfinder quadgramcollocationfinder natural language toolkit collocations and association measures c 2001 2023 nltk project joel nothman jnothman student usyd edu au url https www nltk org for license information see license txt tools to identify collocations words that often appear consecutively within corpora they may also be used to find other associations between word occurrences see manning and schutze ch 5 at https nlp stanford edu fsnlp promo colloc pdf and the text nsp perl package at http ngram sourceforge net finding collocations requires first calculating the frequencies of words and their appearance in the context of other words often the collection of words will then requiring filtering to only retain useful content terms each ngram of words may then be scored according to some association measure in order to determine the relative likelihood of each ngram being a collocation the bigramcollocationfinder and trigramcollocationfinder classes provide these functionalities dependent on being provided a function which scores a ngram given appropriate frequency counts a number of standard association measures are provided in bigram_measures and trigram_measures possible todos consider the distinction between f x _ and f x and whether our approximation is good enough for fragmented data and mention it add a n gram collocation finder with measures which only utilise n gram and unigram counts raw_freq pmi student_t these two unused imports are referenced in collocations doctest an abstract base class for collocation finders whose purpose is to collect collocation candidate frequencies filter and rank them as a minimum collocation finders require the frequencies of each word in a corpus and the joint frequency of word tuples this data should be provided through nltk probability freqdist objects or an identical interface pad the document with the place holder according to the window_size constructs a collocation finder given a collection of documents each of which is a list or iterable of tokens return cls from_words _itertools chain documents generic filter removes ngrams from the frequency distribution if the function returns true when passed an ngram tuple removes candidate ngrams which have frequency less than min_freq removes candidate ngrams w1 w2 where fn w1 w2 evaluates to true removes candidate ngrams w1 w2 where any of fn w1 fn w2 evaluates to true generates of ngram score pairs as determined by the scoring function provided returns a sequence of ngram score pairs ordered from highest to lowest score as determined by the scoring function provided returns the top n ngrams when scored by the given function returns a sequence of ngrams ordered by decreasing score whose scores each exceed the given minimum score a tool for the finding and ranking of bigram collocations or other association measures it is often useful to use from_words rather than constructing an instance directly construct a bigramcollocationfinder given freqdists for appearances of words and possibly non contiguous bigrams construct a bigramcollocationfinder for all bigrams in the given sequence when window_size 2 count non contiguous bigrams in the style of church and hanks s 1990 association ratio returns the score for a given bigram using the given scoring function following church and hanks 1990 counts are scaled by a factor of 1 window_size 1 a tool for the finding and ranking of trigram collocations or other association measures it is often useful to use from_words rather than constructing an instance directly construct a trigramcollocationfinder given freqdists for appearances of words bigrams two words with any word between them and trigrams construct a trigramcollocationfinder for all trigrams in the given sequence constructs a bigram collocation finder with the bigram and unigram data from this finder note that this does not include any filtering applied to this finder returns the score for a given trigram using the given scoring function a tool for the finding and ranking of quadgram collocations or other association measures it is often useful to use from_words rather than constructing an instance directly construct a quadgramcollocationfinder given freqdists for appearances of words bigrams trigrams two words with one word and two words between them three words with a word between them in both variations finds bigram collocations in the files of the webtext corpus slows down loading too much bigram_measures bigramassocmeasures trigram_measures trigramassocmeasures
import itertools as _itertools


from nltk.metrics import (
    BigramAssocMeasures,
    ContingencyMeasures,
    QuadgramAssocMeasures,
    TrigramAssocMeasures,
)
from nltk.metrics.spearman import ranks_from_scores, spearman_correlation
from nltk.probability import FreqDist
from nltk.util import ngrams


class AbstractCollocationFinder:
    

    def __init__(self, word_fd, ngram_fd):
        self.word_fd = word_fd
        self.N = word_fd.N()
        self.ngram_fd = ngram_fd

    @classmethod
    def _build_new_documents(
        cls, documents, window_size, pad_left=False, pad_right=False, pad_symbol=None
    ):
        
        padding = (pad_symbol,) * (window_size - 1)
        if pad_right:
            return _itertools.chain.from_iterable(
                _itertools.chain(doc, padding) for doc in documents
            )
        if pad_left:
            return _itertools.chain.from_iterable(
                _itertools.chain(padding, doc) for doc in documents
            )

    @classmethod
    def from_documents(cls, documents):
        
        
        return cls.from_words(
            cls._build_new_documents(documents, cls.default_ws, pad_right=True)
        )

    @staticmethod
    def _ngram_freqdist(words, n):
        return FreqDist(tuple(words[i : i + n]) for i in range(len(words) - 1))

    def _apply_filter(self, fn=lambda ngram, freq: False):
        
        tmp_ngram = FreqDist()
        for ngram, freq in self.ngram_fd.items():
            if not fn(ngram, freq):
                tmp_ngram[ngram] = freq
        self.ngram_fd = tmp_ngram

    def apply_freq_filter(self, min_freq):
        
        self._apply_filter(lambda ng, freq: freq < min_freq)

    def apply_ngram_filter(self, fn):
        
        self._apply_filter(lambda ng, f: fn(*ng))

    def apply_word_filter(self, fn):
        
        self._apply_filter(lambda ng, f: any(fn(w) for w in ng))

    def _score_ngrams(self, score_fn):
        
        for tup in self.ngram_fd:
            score = self.score_ngram(score_fn, *tup)
            if score is not None:
                yield tup, score

    def score_ngrams(self, score_fn):
        
        return sorted(self._score_ngrams(score_fn), key=lambda t: (-t[1], t[0]))

    def nbest(self, score_fn, n):
        
        return [p for p, s in self.score_ngrams(score_fn)[:n]]

    def above_score(self, score_fn, min_score):
        
        for ngram, score in self.score_ngrams(score_fn):
            if score > min_score:
                yield ngram
            else:
                break


class BigramCollocationFinder(AbstractCollocationFinder):
    

    default_ws = 2

    def __init__(self, word_fd, bigram_fd, window_size=2):
        
        AbstractCollocationFinder.__init__(self, word_fd, bigram_fd)
        self.window_size = window_size

    @classmethod
    def from_words(cls, words, window_size=2):
        
        wfd = FreqDist()
        bfd = FreqDist()

        if window_size < 2:
            raise ValueError("Specify window_size at least 2")

        for window in ngrams(words, window_size, pad_right=True):
            w1 = window[0]
            if w1 is None:
                continue
            wfd[w1] += 1
            for w2 in window[1:]:
                if w2 is not None:
                    bfd[(w1, w2)] += 1
        return cls(wfd, bfd, window_size=window_size)

    def score_ngram(self, score_fn, w1, w2):
        
        n_all = self.N
        n_ii = self.ngram_fd[(w1, w2)] / (self.window_size - 1.0)
        if not n_ii:
            return
        n_ix = self.word_fd[w1]
        n_xi = self.word_fd[w2]
        return score_fn(n_ii, (n_ix, n_xi), n_all)


class TrigramCollocationFinder(AbstractCollocationFinder):
    

    default_ws = 3

    def __init__(self, word_fd, bigram_fd, wildcard_fd, trigram_fd):
        
        AbstractCollocationFinder.__init__(self, word_fd, trigram_fd)
        self.wildcard_fd = wildcard_fd
        self.bigram_fd = bigram_fd

    @classmethod
    def from_words(cls, words, window_size=3):
        
        if window_size < 3:
            raise ValueError("Specify window_size at least 3")

        wfd = FreqDist()
        wildfd = FreqDist()
        bfd = FreqDist()
        tfd = FreqDist()
        for window in ngrams(words, window_size, pad_right=True):
            w1 = window[0]
            if w1 is None:
                continue
            for w2, w3 in _itertools.combinations(window[1:], 2):
                wfd[w1] += 1
                if w2 is None:
                    continue
                bfd[(w1, w2)] += 1
                if w3 is None:
                    continue
                wildfd[(w1, w3)] += 1
                tfd[(w1, w2, w3)] += 1
        return cls(wfd, bfd, wildfd, tfd)

    def bigram_finder(self):
        
        return BigramCollocationFinder(self.word_fd, self.bigram_fd)

    def score_ngram(self, score_fn, w1, w2, w3):
        
        n_all = self.N
        n_iii = self.ngram_fd[(w1, w2, w3)]
        if not n_iii:
            return
        n_iix = self.bigram_fd[(w1, w2)]
        n_ixi = self.wildcard_fd[(w1, w3)]
        n_xii = self.bigram_fd[(w2, w3)]
        n_ixx = self.word_fd[w1]
        n_xix = self.word_fd[w2]
        n_xxi = self.word_fd[w3]
        return score_fn(n_iii, (n_iix, n_ixi, n_xii), (n_ixx, n_xix, n_xxi), n_all)


class QuadgramCollocationFinder(AbstractCollocationFinder):
    

    default_ws = 4

    def __init__(self, word_fd, quadgram_fd, ii, iii, ixi, ixxi, iixi, ixii):
        
        AbstractCollocationFinder.__init__(self, word_fd, quadgram_fd)
        self.iii = iii
        self.ii = ii
        self.ixi = ixi
        self.ixxi = ixxi
        self.iixi = iixi
        self.ixii = ixii

    @classmethod
    def from_words(cls, words, window_size=4):
        if window_size < 4:
            raise ValueError("Specify window_size at least 4")
        ixxx = FreqDist()
        iiii = FreqDist()
        ii = FreqDist()
        iii = FreqDist()
        ixi = FreqDist()
        ixxi = FreqDist()
        iixi = FreqDist()
        ixii = FreqDist()

        for window in ngrams(words, window_size, pad_right=True):
            w1 = window[0]
            if w1 is None:
                continue
            for w2, w3, w4 in _itertools.combinations(window[1:], 3):
                ixxx[w1] += 1
                if w2 is None:
                    continue
                ii[(w1, w2)] += 1
                if w3 is None:
                    continue
                iii[(w1, w2, w3)] += 1
                ixi[(w1, w3)] += 1
                if w4 is None:
                    continue
                iiii[(w1, w2, w3, w4)] += 1
                ixxi[(w1, w4)] += 1
                ixii[(w1, w3, w4)] += 1
                iixi[(w1, w2, w4)] += 1

        return cls(ixxx, iiii, ii, iii, ixi, ixxi, iixi, ixii)

    def score_ngram(self, score_fn, w1, w2, w3, w4):
        n_all = self.N
        n_iiii = self.ngram_fd[(w1, w2, w3, w4)]
        if not n_iiii:
            return
        n_iiix = self.iii[(w1, w2, w3)]
        n_xiii = self.iii[(w2, w3, w4)]
        n_iixi = self.iixi[(w1, w2, w4)]
        n_ixii = self.ixii[(w1, w3, w4)]

        n_iixx = self.ii[(w1, w2)]
        n_xxii = self.ii[(w3, w4)]
        n_xiix = self.ii[(w2, w3)]
        n_ixix = self.ixi[(w1, w3)]
        n_ixxi = self.ixxi[(w1, w4)]
        n_xixi = self.ixi[(w2, w4)]

        n_ixxx = self.word_fd[w1]
        n_xixx = self.word_fd[w2]
        n_xxix = self.word_fd[w3]
        n_xxxi = self.word_fd[w4]
        return score_fn(
            n_iiii,
            (n_iiix, n_iixi, n_ixii, n_xiii),
            (n_iixx, n_ixix, n_ixxi, n_xixi, n_xxii, n_xiix),
            (n_ixxx, n_xixx, n_xxix, n_xxxi),
            n_all,
        )


def demo(scorer=None, compare_scorer=None):
    
    from nltk.metrics import (
        BigramAssocMeasures,
        ranks_from_scores,
        spearman_correlation,
    )

    if scorer is None:
        scorer = BigramAssocMeasures.likelihood_ratio
    if compare_scorer is None:
        compare_scorer = BigramAssocMeasures.raw_freq

    from nltk.corpus import stopwords, webtext

    ignored_words = stopwords.words("english")
    word_filter = lambda w: len(w) < 3 or w.lower() in ignored_words

    for file in webtext.fileids():
        words = [word.lower() for word in webtext.words(file)]

        cf = BigramCollocationFinder.from_words(words)
        cf.apply_freq_filter(3)
        cf.apply_word_filter(word_filter)

        corr = spearman_correlation(
            ranks_from_scores(cf.score_ngrams(scorer)),
            ranks_from_scores(cf.score_ngrams(compare_scorer)),
        )
        print(file)
        print("\t", [" ".join(tup) for tup in cf.nbest(scorer, 15)])
        print(f"\t Correlation to {compare_scorer.__name__}: {corr:0.4f}")






if __name__ == "__main__":
    import sys

    from nltk.metrics import BigramAssocMeasures

    try:
        scorer = eval("BigramAssocMeasures." + sys.argv[1])
    except IndexError:
        scorer = None
    try:
        compare_scorer = eval("BigramAssocMeasures." + sys.argv[2])
    except IndexError:
        compare_scorer = None

    demo(scorer, compare_scorer)

__all__ = [
    "BigramCollocationFinder",
    "TrigramCollocationFinder",
    "QuadgramCollocationFinder",
]