# Words by the Numbers 4, Exercise 2: raw counts against PMI, with and without common words.
import math, re
from collections import Counter
from pydoc_data.topics import topics

def tokens(text):
    return re.findall(r"[a-z]+(?:'[a-z]+)?", text.lower())

uni, bi = Counter(), Counter()
for doc in topics.values():
    w = tokens(doc)
    uni.update(w)
    bi.update(zip(w, w[1:]))
N = sum(uni.values())
SMALL = {"the", "of", "in", "to", "a", "is", "and", "be", "it", "if", "as", "for", "an", "by", "or", "on", "that", "this", "with", "are", "can", "not"}

def pmi(pair):
    return math.log2(bi[pair] / (uni[pair[0]] * uni[pair[1]] / N))

def show(label, pairs):
    print(label, [(" ".join(p), bi[p], round(pmi(p), 1)) for p in pairs])

by_count = [p for p, _ in bi.most_common(6)]
no_small = [p for p, _ in bi.most_common() if not (set(p) & SMALL)][:6]
by_pmi = sorted((p for p in bi if bi[p] >= 10), key=pmi, reverse=True)[:6]
show("raw count:           ", by_count)
show("count, no small words:", no_small)
show("PMI, count >= 10:     ", by_pmi)

# Output when run (Python 3.14; the documentation-text figures depend on the Python version):
# raw count:            [('of the', 532, 2.3), ('in the', 447, 2.0), ('to the', 225, 1.2), ('if the', 209, 1.9), ('in version', 205, 5.4), ('can be', 193, 5.9)]
# count, no small words: [('self other', 94, 6.9), ('other object', 81, 5.3), ('x y', 73, 6.3), ('see also', 67, 6.7), ('e g', 51, 8.6), ('except clause', 48, 7.4)]
# PMI, count >= 10:      [('auditing event', 11, 12.4), ('square brackets', 16, 12.2), ('hand side', 14, 12.0), ('nearest enclosing', 11, 11.6), ('simple terms', 11, 11.4), ('codeobject co', 19, 11.4)]