Grammars and Constituency Structure › Code Implementation

import math
from dataclasses import dataclass, field
from typing import Dict, List, Tuple, Union


@dataclass
class Tree:
    label: str
    children: List[Union["Tree", str]] = field(default_factory=list)

    def leaf_tokens(self) -> List[str]:
        tokens = []
        for child in self.children:
            if isinstance(child, Tree):
                tokens.extend(child.leaf_tokens())
            else:
                tokens.append(child)
        return tokens

    def pretty(self, indent: int = 0) -> str:
        pad = "  " * indent
        lines = [f"{pad}{self.label}"]
        for child in self.children:
            if isinstance(child, Tree):
                lines.append(child.pretty(indent + 1))
            else:
                lines.append(f"{'  ' * (indent + 1)}{child}")
        return "\n".join(lines)

    def bracket(self) -> str:
        parts = []
        for child in self.children:
            parts.append(child.bracket() if isinstance(child, Tree) else child)
        return f"({self.label} {' '.join(parts)})"

Rule = Tuple[str, Tuple[str, ...]]

rules: List[Rule] = [
    ("S", ("NP", "VP")),
    ("VP", ("V", "NP")),
    ("VP", ("VP", "PP")),
    ("NP", ("Det", "N")),
    ("NP", ("NP", "PP")),
    ("PP", ("P", "NP")),
    ("Det", ("the",)),
    ("Det", ("a",)),
    ("N", ("dog",)),
    ("N", ("cat",)),
    ("N", ("telescope",)),
    ("V", ("chased",)),
    ("P", ("with",)),
]

def cky_parse(words: List[str], grammar: List[Rule]):
    n = len(words)
    chart = [[{} for _ in range(n + 1)] for _ in range(n + 1)]

    nonterminals = {lhs for lhs, _ in grammar}
    lexical: Dict[str, List[str]] = {}
    binary: Dict[Tuple[str, str], List[str]] = {}

    for lhs, rhs in grammar:
        if len(rhs) == 1 and rhs[0] not in nonterminals:
            lexical.setdefault(rhs[0], []).append(lhs)
        elif len(rhs) == 2:
            binary.setdefault((rhs[0], rhs[1]), []).append(lhs)

    for i, word in enumerate(words):
        for lhs in lexical.get(word, []):
            chart[i][i + 1].setdefault(lhs, []).append(Tree(lhs, [word]))

    for span in range(2, n + 1):
        for i in range(n - span + 1):
            j = i + span
            cell = chart[i][j]
            for k in range(i + 1, j):
                left = chart[i][k]
                right = chart[k][j]
                for (b, c), lhs_list in binary.items():
                    if b not in left or c not in right:
                        continue
                    for left_tree in left[b]:
                        for right_tree in right[c]:
                            for a in lhs_list:
                                cell.setdefault(a, []).append(
                                    Tree(a, [left_tree, right_tree])
                                )

    return chart, chart[0][n].get("S", [])

sentence = "the dog chased the cat with a telescope"
tokens = sentence.split()
chart, parses = cky_parse(tokens, rules)

rule_probs: Dict[Rule, float] = {
    ("S", ("NP", "VP")): 1.0,
    ("VP", ("V", "NP")): 0.7,
    ("VP", ("VP", "PP")): 0.3,
    ("NP", ("Det", "N")): 0.8,
    ("NP", ("NP", "PP")): 0.2,
    ("PP", ("P", "NP")): 1.0,
    ("Det", ("the",)): 0.8,
    ("Det", ("a",)): 0.2,
    ("N", ("dog",)): 0.4,
    ("N", ("cat",)): 0.4,
    ("N", ("telescope",)): 0.2,
    ("V", ("chased",)): 1.0,
    ("P", ("with",)): 1.0,
}


def tree_probability(tree: Tree, probs: Dict[Rule, float]) -> float:
    rhs = tuple(
        child.label if isinstance(child, Tree) else child
        for child in tree.children
    )
    probability = probs[(tree.label, rhs)]
    for child in tree.children:
        if isinstance(child, Tree):
            probability *= tree_probability(child, probs)
    return probability

import numpy as np

parse_prob_values = np.array([prob for prob, _ in scored])
parse_attachment_labels = [
    "VP attachment" if "(VP (VP" in tree.bracket() else "NP attachment"
    for _, tree in scored
]

def catalan(n: int) -> int:
    return math.comb(2 * n, n) // (n + 1)