Probabilistic CKY in Python
CYK variant computing the highest-probability parse of a string under a PCFG in cubic time.
- File
probabilistic_cky.py, 196 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/probabilistic-cky/. The authors are listed in prov/W4N6QS.jsonl.
"""Probabilistic (Viterbi) CKY for a probabilistic grammar in Chomsky normal form.
A grammar is a list of rules (A, rhs, q), where rhs is a tuple of symbols and q
is a nonnegative number, the probability of the production A -> rhs. The
nonterminals are the left sides. Every rule must be A -> B C with B, C
nonterminals, A -> a with a terminal, or start -> () when the start symbol is on
no right side. A parse tree is (A, children), where each child is a tree or a
terminal symbol.
The probabilities can be Fractions (exact), floats, or, with log_space=True,
floats that are added as logarithms so that long inputs do not underflow.
"""
import math
from typing import Sequence, Union
Tree = tuple[str, tuple[Union["Tree", str], ...]]
Rule = tuple[str, tuple[str, ...], object]
def check_cnf(rules: Sequence[Rule], start: str) -> set[str]:
"""Raise ValueError unless the rules are a weighted grammar in Chomsky normal form."""
nts = {a for a, _, _ in rules}
has_empty = False
for a, rhs, q in rules:
if q < 0:
raise ValueError(f"negative probability for {a} -> {rhs}")
if len(rhs) == 2 and rhs[0] in nts and rhs[1] in nts:
continue
if len(rhs) == 1 and rhs[0] not in nts:
continue
if len(rhs) == 0 and a == start:
has_empty = True
continue
raise ValueError(f"not in Chomsky normal form: {a} -> {rhs}")
if has_empty and any(start in rhs for _, rhs, _ in rules):
raise ValueError("the start symbol has an empty production and is on a right side")
return nts
def _chart(rules, w, log_space):
"""Fill the Viterbi chart. best[i][j] maps A to the largest probability of a
tree with root A and yield w[i..j] (0-based, both ends included); back holds
the split (k, B, C), or the terminal for j = i, of one best tree."""
if log_space:
lift = [(a, rhs, math.log(q)) for a, rhs, q in rules if q > 0]
else:
lift = list(rules)
n = len(w)
lexical: dict[str, list[tuple[str, object]]] = {}
binary = []
for a, rhs, q in lift:
if len(rhs) == 1:
lexical.setdefault(rhs[0], []).append((a, q))
elif len(rhs) == 2:
binary.append((a, rhs[0], rhs[1], q))
best: list[list[dict[str, object]]] = [[{} for _ in range(n)] for _ in range(n)]
back: dict[tuple[int, int, str], object] = {}
for i in range(n): # @line L2
cell = best[i][i]
for a, q in lexical.get(w[i], ()):
if a not in cell or q > cell[a]:
cell[a] = q # @line L3
back[(i, i, a)] = w[i]
for length in range(2, n + 1): # @line L4
for i in range(n - length + 1): # @line L5
j = i + length - 1 # @line L6
cell = best[i][j]
for k in range(i, j): # @line L7
left, right = best[i][k], best[k + 1][j]
if not left or not right:
continue
for a, b, c, q in binary: # @line L8
if b in left and c in right:
if log_space:
p = q + left[b] + right[c] # @line L9
else:
p = q * left[b] * right[c] # @line L9
if a not in cell or p > cell[a]: # @line L10
cell[a] = p
back[(i, j, a)] = (k, b, c)
return best, back
def _build_tree(back, start: str, n: int) -> Tree:
"""Follow the back pointers from (0, n - 1, start) with an explicit stack."""
done: dict[tuple[int, int, str], Tree] = {}
stack = [(0, n - 1, start, False)]
while stack:
i, j, a, expanded = stack.pop()
b = back[(i, j, a)]
if i == j:
done[(i, j, a)] = (a, (b,))
continue
k, left, right = b
kids = ((i, k, left), (k + 1, j, right))
if expanded:
done[(i, j, a)] = (a, (done[kids[0]], done[kids[1]]))
else:
stack.append((i, j, a, True))
for key in kids:
stack.append((key[0], key[1], key[2], False))
return done[(0, n - 1, start)]
def viterbi_cky(rules: Sequence[Rule], start: str, w: Sequence[str], log_space: bool = False):
"""Return (p, tree): the largest probability p of a parse tree of w and one tree
that has it. When w has no parse tree, return (0, None), or (-inf, None) when
log_space is True, in which case p is the natural logarithm of the probability.
Time Theta(n^3 g) for n = len(w) and g rules, space Theta(n^2 v).
"""
check_cnf(rules, start)
zero = -math.inf if log_space else 0
if len(w) == 0: # @line L1
empty = [q for a, rhs, q in rules if a == start and len(rhs) == 0 and q > 0]
if not empty:
return zero, None
q = max(empty)
return (math.log(q) if log_space else q), (start, ())
best, back = _chart(rules, w, log_space)
top = best[0][len(w) - 1]
if start not in top:
return zero, None
return top[start], _build_tree(back, start, len(w)) # @line L11
def best_chart(rules: Sequence[Rule], start: str, w: Sequence[str]) -> list[list[dict[str, object]]]:
"""Return the chart: entry [i][j][A] is the best probability of A over w[i..j]."""
check_cnf(rules, start)
return _chart(rules, w, False)[0]
def inside_probability(rules: Sequence[Rule], start: str, w: Sequence[str]):
"""Return P(w), the sum of the probabilities of all parse trees of w.
The same loops as viterbi_cky, with the maximum replaced by a sum.
"""
check_cnf(rules, start)
n = len(w)
if n == 0:
return sum((q for a, rhs, q in rules if a == start and len(rhs) == 0), 0)
inside: list[list[dict[str, object]]] = [[{} for _ in range(n)] for _ in range(n)]
for i in range(n):
cell = inside[i][i]
for a, rhs, q in rules:
if len(rhs) == 1 and rhs[0] == w[i]:
cell[a] = cell.get(a, 0) + q
binary = [(a, rhs[0], rhs[1], q) for a, rhs, q in rules if len(rhs) == 2]
for length in range(2, n + 1):
for i in range(n - length + 1):
j = i + length - 1
cell = inside[i][j]
for k in range(i, j):
left, right = inside[i][k], inside[k + 1][j]
if not left or not right:
continue
for a, b, c, q in binary:
if b in left and c in right:
cell[a] = cell.get(a, 0) + q * left[b] * right[c]
return inside[0][n - 1].get(start, 0)
def tree_probability(rules: Sequence[Rule], tree: Tree):
"""Return the product of the probabilities of the rules used in the tree.
If a production appears several times in the list, its largest probability is used.
"""
q: dict[tuple[str, tuple[str, ...]], object] = {}
for a, rhs, p in rules:
if (a, rhs) not in q or p > q[(a, rhs)]:
q[(a, rhs)] = p
total = 1
stack = [tree]
while stack:
a, kids = stack.pop()
rhs = tuple(k if isinstance(k, str) else k[0] for k in kids)
total = total * q[(a, rhs)]
stack.extend(k for k in kids if not isinstance(k, str))
return total
def tree_yield(tree: Tree) -> list[str]:
"""Return the terminal symbols at the leaves of the tree, from left to right."""
out = []
stack: list[object] = [tree]
while stack:
t = stack.pop()
if isinstance(t, str):
out.append(t)
else:
stack.extend(reversed(t[1]))
return out
test_probabilistic_cky.py 235 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/probabilistic-cky/. The authors are listed in prov/W4N6QS.jsonl.
import itertools
import math
import random
from fractions import Fraction as F
import pytest
from probabilistic_cky import (
best_chart,
check_cnf,
inside_probability,
tree_probability,
tree_yield,
viterbi_cky,
)
ASTRO = [
("S", ("NP", "VP"), F(1)),
("PP", ("P", "NP"), F(1)),
("VP", ("V", "NP"), F(7, 10)),
("VP", ("VP", "PP"), F(3, 10)),
("P", ("with",), F(1)),
("V", ("saw",), F(1)),
("NP", ("NP", "PP"), F(4, 10)),
("NP", ("astronomers",), F(1, 10)),
("NP", ("ears",), F(18, 100)),
("NP", ("saw",), F(4, 100)),
("NP", ("stars",), F(18, 100)),
("NP", ("telescopes",), F(1, 10)),
]
def all_trees(rules, start, w):
"""Brute force: enumerate every leftmost derivation of w, with an explicit stack.
In Chomsky normal form no symbol other than an empty start derives the empty
string, so a sentential form longer than w can be dropped. Returns the list
of the probabilities of the derivations, one per parse tree.
"""
nts = {a for a, _, _ in rules}
by_lhs = {}
for a, rhs, q in rules:
by_lhs.setdefault(a, []).append((rhs, q))
n = len(w)
out = []
stack = [((start,), 1)]
while stack:
form, p = stack.pop()
k = 0
while k < len(form) and form[k] not in nts:
k += 1
if tuple(form[:k]) != tuple(w[:k]):
continue
if k == len(form):
if len(form) == n:
out.append(p)
continue
for rhs, q in by_lhs[form[k]]:
new = form[:k] + rhs + form[k + 1:]
if len(new) <= n or (n == 0 and len(new) == 0):
stack.append((new, p * q))
return out
def random_grammar(rng):
"""A random weighted grammar in CNF over the nonterminals S, A, B, C and the terminals a, b.
It may be ambiguous, left or right recursive, have S on right sides, or have
S -> epsilon (only when S is on no right side). The weights are proper
probabilities most of the time and arbitrary nonnegative weights otherwise.
"""
nts = ["S", "A", "B", "C"][: rng.randint(1, 4)]
allow_s_rhs = rng.random() < 0.5
rhs_nts = nts if allow_s_rhs else [x for x in nts if x != "S"] or ["S"]
if not allow_s_rhs and rhs_nts == ["S"]:
allow_s_rhs = True
shapes = []
for a in nts:
for t in "ab":
if rng.random() < 0.5:
shapes.append((a, (t,)))
for b, c in itertools.product(rhs_nts, repeat=2):
if rng.random() < 0.3:
shapes.append((a, (b, c)))
if not allow_s_rhs and rng.random() < 0.4:
shapes.append(("S", ()))
if rng.random() < 0.2 and shapes:
shapes.append(rng.choice(shapes)) # a duplicate production
proper = rng.random() < 0.8
weights = [F(rng.randint(0, 6), 1) if not proper else F(rng.randint(1, 6)) for _ in shapes]
rules = []
if proper:
total = {}
for (a, _), q in zip(shapes, weights):
total[a] = total.get(a, 0) + q
for (a, rhs), q in zip(shapes, weights):
rules.append((a, rhs, q / total[a]))
else:
for (a, rhs), q in zip(shapes, weights):
rules.append((a, rhs, q / 3))
# Make sure every right-side nonterminal is a left side, as CNF requires here.
lhs = {a for a, _, _ in rules}
return [r for r in rules if all(x in lhs or x in "ab" for x in r[1])]
def check_answer(rules, start, w, got_p, got_tree, trees):
if not trees or max(trees) == 0:
assert got_tree is None or got_p == 0
return
want = max(trees)
assert got_p == want
assert got_tree is not None
assert got_tree[0] == start
assert tree_yield(got_tree) == list(w)
assert tree_probability(rules, got_tree) == want
def test_random_grammars_against_enumeration():
rng = random.Random(20261106)
comparisons = 0
nonempty = 0
while comparisons < 3000:
rules = random_grammar(rng)
try:
check_cnf(rules, "S")
except ValueError:
continue
if not any(a == "S" for a, _, _ in rules):
continue
for n in range(0, 6):
for w in itertools.product("ab", repeat=n):
if n >= 4 and rng.random() < 0.6:
continue
trees = all_trees(rules, "S", w)
p, tree = viterbi_cky(rules, "S", w)
check_answer(rules, "S", w, p, tree, trees)
assert inside_probability(rules, "S", w) == sum(trees, F(0))
if trees and max(trees) > 0:
nonempty += 1
lp, ltree = viterbi_cky(rules, "S", w, log_space=True)
assert math.isclose(lp, math.log(max(trees)), rel_tol=1e-9, abs_tol=1e-9)
assert tree_yield(ltree) == list(w)
comparisons += 1
assert nonempty > 200
def test_astronomers_example():
w = "astronomers saw stars with ears".split()
p, tree = viterbi_cky(ASTRO, "S", w)
assert p == F(9072, 10**7)
assert inside_probability(ASTRO, "S", w) == F(15876, 10**7)
assert p / inside_probability(ASTRO, "S", w) == F(4, 7)
# The best tree attaches "with ears" to "stars".
vp = tree[1][1]
assert vp[0] == "VP" and vp[1][0][0] == "V" and vp[1][1][0] == "NP"
assert vp[1][1][1][0][0] == "NP" and vp[1][1][1][1][0] == "PP"
chart = best_chart(ASTRO, "S", w)
assert chart[1][2] == {"VP": F(126, 1000)}
assert chart[3][4] == {"PP": F(18, 100)}
assert chart[0][2] == {"S": F(126, 10**4)}
assert chart[2][4] == {"NP": F(1296, 10**5)}
assert chart[1][4] == {"VP": F(9072, 10**6)}
assert chart[0][1] == {} and chart[0][3] == {}
assert chart[0][4] == {"S": F(9072, 10**7)}
assert len(all_trees(ASTRO, "S", w)) == 2
def test_floats_match_fractions():
floats = [(a, r, float(q)) for a, r, q in ASTRO]
w = "astronomers saw stars with ears".split()
p, _ = viterbi_cky(floats, "S", w)
assert math.isclose(p, 0.0009072)
def test_empty_input():
g = [("S", (), F(1, 4)), ("S", ("A", "A"), F(3, 4)), ("A", ("a",), F(1))]
assert viterbi_cky(g, "S", []) == (F(1, 4), ("S", ()))
assert inside_probability(g, "S", []) == F(1, 4)
g2 = [("S", ("a",), F(1))]
assert viterbi_cky(g2, "S", []) == (0, None)
assert viterbi_cky(g2, "S", [], log_space=True) == (-math.inf, None)
assert inside_probability(g2, "S", []) == 0
def test_no_parse():
g = [("S", ("A", "B"), F(1)), ("A", ("a",), F(1)), ("B", ("b",), F(1))]
assert viterbi_cky(g, "S", list("ba")) == (0, None)
assert viterbi_cky(g, "S", list("abc")) == (0, None)
assert viterbi_cky(g, "S", list("ab")) == (F(1), ("S", (("A", ("a",)), ("B", ("b",)))))
def test_rejects_non_cnf():
bad = [
[("S", ("A",), F(1)), ("A", ("a",), F(1))], # unit production
[("S", ("A", "A"), F(1)), ("A", (), F(1))], # empty production of a non-start symbol
[("S", ("a", "S"), F(1, 2)), ("S", ("a",), F(1, 2))], # terminal in a binary rule
[("S", ("S", "S"), F(1, 2)), ("S", (), F(1, 2))], # S -> epsilon with S on a right side
[("S", ("a",), F(-1))], # negative probability
[("S", ("A", "B", "C"), F(1))], # long right side
]
for g in bad:
with pytest.raises(ValueError):
viterbi_cky(g, "S", ["a"])
def test_ties_and_ambiguity():
# S -> S S | a with equal halves: every binary tree of a^n has the same probability.
g = [("S", ("S", "S"), F(1, 2)), ("S", ("a",), F(1, 2))]
catalan = [1, 1, 2, 5, 14, 42]
for n in range(1, 7):
p, tree = viterbi_cky(g, "S", ["a"] * n)
assert p == F(1, 2) ** (2 * n - 1)
assert inside_probability(g, "S", ["a"] * n) == catalan[n - 1] * p
assert tree_yield(tree) == ["a"] * n
def test_log_space_avoids_underflow():
# A right-branching grammar: S -> A S | a with tiny lexical probabilities.
g = [("S", ("A", "S"), 0.5), ("S", ("a",), 0.5), ("A", ("a",), 1e-3)]
n = 150
w = ["a"] * n
p, _ = viterbi_cky(g, "S", w)
assert p == 0.0 # the product 0.5^150 * 1e-447 underflows a double
lp, tree = viterbi_cky(g, "S", w, log_space=True)
want = (n - 1) * (math.log(0.5) + math.log(1e-3)) + math.log(0.5)
assert math.isclose(lp, want, rel_tol=1e-12)
depth, t = 0, tree
while len(t[1]) == 2:
t = t[1][1]
depth += 1
assert depth == n - 1
assert tree_yield(tree) == w
What the next level needs
- the tests did not run in this build