Constituency and Dependency Parsing › Code Implementation

from __future__ import annotations

from collections import deque
from itertools import product

import numpy as np

lexicon = {
    "Det": {"the", "a"},
    "N": {"man", "dog", "telescope", "park"},
    "V": {"saw", "watched"},
    "P": {"with", "in"},
}

binary_rules = [
    ("NP", "Det", "N"),
    ("VP", "V", "NP"),
    ("PP", "P", "NP"),
    ("NP", "NP", "PP"),
    ("VP", "VP", "PP"),
    ("S", "NP", "VP"),
]

def cky_parse(tokens):
    n = len(tokens)
    chart = [[{} for _ in range(n + 1)] for _ in range(n + 1)]

    for i, token in enumerate(tokens):
        for nt, words in lexicon.items():
            if token in words:
                chart[i][i + 1][nt] = [(nt, token)]

    for length in range(2, n + 1):
        for i in range(n - length + 1):
            j = i + length
            cell = {}
            for k in range(i + 1, j):
                left = chart[i][k]
                right = chart[k][j]
                for parent, left_nt, right_nt in binary_rules:
                    if left_nt in left and right_nt in right:
                        for left_tree in left[left_nt]:
                            for right_tree in right[right_nt]:
                                cell.setdefault(parent, []).append(
                                    (parent, left_tree, right_tree)
                                )
            chart[i][j] = cell

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

def render_tree(tree, indent=0):
    if len(tree) == 2:
        label, word = tree
        return "  " * indent + f"{label} -> {word}"
    label, left, right = tree
    lines = ["  " * indent + label]
    lines.append(render_tree(left, indent + 1))
    lines.append(render_tree(right, indent + 1))
    return "\n".join(lines)

tokens = "the man saw the dog with a telescope".split()
chart, parses = cky_parse(tokens)

def chart_rows(chart, tokens):
    rows = []
    n = len(tokens)
    for length in range(1, n + 1):
        for i in range(n - length + 1):
            j = i + length
            nts = sorted(chart[i][j])
            if nts:
                rows.append((i, j, " ".join(tokens[i:j]), ", ".join(nts)))
    return rows


rows = chart_rows(chart, tokens)

def apply_transition(stack, buffer, arcs, action, label=None):
    if action == "SHIFT":
        if not buffer:
            raise ValueError("SHIFT requires a non-empty buffer")
        stack.append(buffer.popleft())
    elif action == "LEFT-ARC":
        if len(stack) < 2 or stack[-2] == 0:
            raise ValueError(
                "LEFT-ARC requires two items and cannot attach ROOT"
            )
        head, dep = stack[-1], stack[-2]
        arcs.append((head, dep, label))
        stack.pop(-2)
    elif action == "RIGHT-ARC":
        if len(stack) < 2:
            raise ValueError("RIGHT-ARC requires two stack items")
        if stack[-2] == 0 and (buffer or len(stack) != 2):
            raise ValueError("The ROOT arc must be the final transition")
        head, dep = stack[-2], stack[-1]
        arcs.append((head, dep, label))
        stack.pop()
    else:
        raise ValueError(f"Unknown action: {action}")


def run_transitions(words, transitions):
    stack = [0]
    buffer = deque(range(1, len(words)))
    arcs = []
    for action, label in transitions:
        apply_transition(stack, buffer, arcs, action, label)
    return stack, list(buffer), arcs

words = ["<ROOT>", "She", "saw", "him"]
transitions = [
    ("SHIFT", None),
    ("SHIFT", None),
    ("LEFT-ARC", "nsubj"),
    ("SHIFT", None),
    ("RIGHT-ARC", "dobj"),
    ("RIGHT-ARC", "root"),
]
stack, buffer, arcs = run_transitions(words, transitions)

def is_valid_tree(heads):
    n = len(heads) - 1
    if sum(1 for d in range(1, n + 1) if heads[d] == 0) != 1:
        return False
    for dep in range(1, n + 1):
        seen = set()
        cur = dep
        while cur != 0:
            if cur in seen:
                return False
            seen.add(cur)
            cur = heads[cur]
    return True


def brute_force_mst(n, arc_scores):
    best_total = float("-inf")
    best_heads = None
    for parents in product(range(n + 1), repeat=n):
        heads = [None] + list(parents)
        if any(heads[d] == d for d in range(1, n + 1)):
            continue
        if not is_valid_tree(heads):
            continue
        total = sum(arc_scores.get((heads[d], d), 0.0) for d in range(1, n + 1))
        if total > best_total:
            best_total = total
            best_heads = heads
    return best_heads, best_total

words = ["<ROOT>", "She", "saw", "the", "man"]
n = len(words) - 1
gold_arcs = {(2, 1): "nsubj", (0, 2): "root", (2, 4): "dobj", (4, 3): "det"}

arc_scores = {}
for h in range(n + 1):
    for d in range(1, n + 1):
        if h == d:
            continue
        arc_scores[(h, d)] = 5.0 if (h, d) in gold_arcs else -1.0

best_heads, best_total = brute_force_mst(n, arc_scores)

rng = np.random.default_rng(7)
dim = 8
H = rng.normal(size=(n + 1, dim))
W_head = rng.normal(scale=0.3, size=(dim, dim))
W_dep = rng.normal(scale=0.3, size=(dim, dim))
U = rng.normal(scale=0.3, size=(dim, dim))

H_head = H @ W_head
H_dep = H @ W_dep
b_head = rng.normal(scale=0.3, size=dim)
b_dep = rng.normal(scale=0.3, size=dim)
bias = float(rng.normal(scale=0.3))
raw_scores = (
    H_head @ U @ H_dep.T
    + (H_head @ b_head)[:, None]
    + (H_dep @ b_dep)[None, :]
    + bias
)

neural_scores = {}
for h in range(n + 1):
    for d in range(1, n + 1):
        if h == d:
            continue
        neural_scores[(h, d)] = float(raw_scores[h, d])

neural_heads, neural_total = brute_force_mst(n, neural_scores)