N-gram Language Models

This small script computes bigram MLE and add-one probabilities on a tiny corpus. It stays deliberately small so the count table and sentence probability are inspectable.

from collections import Counter, defaultdict
from math import exp, log, isinf

train = [
    "the cat sat",
    "the cat slept",
    "the cat purred",
    "the dog sat",
    "a dog barked",
    "the dog barked",
]

def tokens(sentence):
    return ["<s>"] + sentence.split() + ["</s>"]

pairs = Counter()
contexts = Counter()
vocab = set(["</s>"])

for sentence in train:
    seq = tokens(sentence)
    vocab.update(seq[1:])
    for left, right in zip(seq, seq[1:]):
        pairs[(left, right)] += 1
        contexts[left] += 1

def prob(left, right, alpha=0.0):
    return (pairs[(left, right)] + alpha) / (contexts[left] + alpha * len(vocab))

def score(sentence, alpha=0.0):
    seq = tokens(sentence)
    logp = 0.0
    zero = []
    for left, right in zip(seq, seq[1:]):
        p = prob(left, right, alpha)
        if p == 0:
            zero.append((left, right))
        else:
            logp += log(p)
    if zero:
        return 0.0, float("inf"), zero
    n = len(seq) - 1
    return exp(logp), exp(-logp / n), zero

for word in ["cat", "dog", "barked", "slept"]:
    print(f"P({word} | the) =", round(prob("the", word), 3))

print("MLE:", score("the dog slept", alpha=0.0))
print("add-one:", score("the dog slept", alpha=1.0))

The unsmoothed model gives the dog slept zero probability because it never saw dog slept. Add-one smoothing gives that unseen bigram a small nonzero probability, so the sentence probability and perplexity become finite.