Semantic Parsing and Meaning Representations › Code Implementation

import re
import sqlite3
from dataclasses import dataclass
from typing import Any

import matplotlib.pyplot as plt
import numpy as np

@dataclass(frozen=True)
class Var:
    name: str


@dataclass(frozen=True)
class Const:
    name: str


@dataclass(frozen=True)
class App:
    fn: Any
    arg: Any


@dataclass(frozen=True)
class Lam:
    var: str
    body: Any

def evaluate(term, env, constants):
    if isinstance(term, Var):
        return env[term.name]
    if isinstance(term, Const):
        return constants[term.name]
    if isinstance(term, Lam):

        def abstraction(arg, _var=term.var, _body=term.body, _env=env):
            new_env = dict(_env)
            new_env[_var] = arg
            return evaluate(_body, new_env, constants)

        return abstraction
    if isinstance(term, App):
        fn = evaluate(term.fn, env, constants)
        arg = evaluate(term.arg, env, constants)
        return fn(arg)
    raise TypeError(f"Unknown term: {term!r}")

def show(term):
    if isinstance(term, Var):
        return term.name
    if isinstance(term, Const):
        return term.name
    if isinstance(term, Lam):
        return f"(λ{term.var}.{show(term.body)})"
    if isinstance(term, App):
        return f"({show(term.fn)} {show(term.arg)})"
    return str(term)

domain = ["alice", "bob", "carol", "b1", "b2"]
students = {"alice", "bob", "carol"}
books = {"b1", "b2"}
reads = {("alice", "b1"), ("bob", "b2"), ("carol", "b1")}
likes = {("alice", "bob"), ("bob", "carol")}

constants = {
    "alice": "alice",
    "bob": "bob",
    "student": lambda x: x in students,
    "book": lambda x: x in books,
    "read": lambda x: lambda y: (x, y) in reads,
    "like": lambda x: lambda y: (x, y) in likes,
    "every": lambda p: lambda q: all((not p(x)) or q(x) for x in domain),
    "a": lambda p: lambda q: any(p(x) and q(x) for x in domain),
}

alice_likes_bob = App(App(Const("like"), Const("alice")), Const("bob"))

every_student = App(Const("every"), Lam("x", App(Const("student"), Var("x"))))
a_book = App(Const("a"), Lam("y", App(Const("book"), Var("y"))))
reads_a_book = Lam(
    "x",
    App(a_book, Lam("y", App(App(Const("read"), Var("x")), Var("y")))),
)
sentence = App(every_student, reads_a_book)

no_bob_reads = {("alice", "b1"), ("carol", "b1")}
constants_without_bob = dict(constants)
constants_without_bob["read"] = lambda x: lambda y: (x, y) in no_bob_reads

@dataclass
class AMRGraph:
    nodes: dict[str, str]
    edges: list[tuple[str, str, str]]


amr = AMRGraph(
    nodes={"w": "want-01", "b": "boy", "g": "go-02"},
    edges=[
        ("w", "ARG0", "b"),
        ("w", "ARG1", "g"),
        ("g", "ARG0", "b"),
    ],
)


def amr_triples(graph):
    return [
        (graph.nodes[src], role, graph.nodes.get(tgt, tgt))
        for src, role, tgt in graph.edges
    ]

conn = sqlite3.connect(":memory:")
cur = conn.cursor()

cur.executescript(
    """
    CREATE TABLE departments (
        department_id INTEGER PRIMARY KEY,
        name TEXT NOT NULL
    );
    CREATE TABLE students (
        student_id INTEGER PRIMARY KEY,
        name TEXT NOT NULL,
        department_id INTEGER REFERENCES departments(department_id)
    );
    CREATE TABLE enrollments (
        student_id INTEGER REFERENCES students(student_id),
        course_id INTEGER NOT NULL
    );

    INSERT INTO departments VALUES
        (1, 'Computer Science'),
        (2, 'Mathematics'),
        (3, 'Physics');

    INSERT INTO students VALUES
        (1, 'Alice', 1),
        (2, 'Bob', 1),
        (3, 'Carol', 2),
        (4, 'Dave', 1);

    INSERT INTO enrollments VALUES
        (1, 101),
        (2, 101),
        (4, 102),
        (3, 201);
    """
)

question = "Which departments have more than two enrolled students?"
number_words = {"one": 1, "two": 2, "three": 3, "four": 4}
match = re.search(r"more than (\w+)", question)
threshold = number_words[match.group(1)]

sql = """
SELECT d.name
FROM departments AS d
JOIN students AS s
  ON s.department_id = d.department_id
JOIN enrollments AS e
  ON e.student_id = s.student_id
GROUP BY d.department_id, d.name
HAVING COUNT(DISTINCT e.student_id) > ?
ORDER BY d.name;
"""

rows = cur.execute(sql, (threshold,)).fetchall()

semantic_tree = {
    "label": "S : t",
    "children": [
        {"label": "Subject : every(student)"},
        {
            "label": "VP : property of x\nx reads a book",
            "children": [
                {"label": "Verb : read"},
                {"label": "Object : a(book)"},
            ],
        },
    ],
}


def leaf_count(node):
    if "children" not in node:
        return 1
    return sum(leaf_count(child) for child in node["children"])


def layout(node, x0, x1, y, positions, edges):
    x = (x0 + x1) / 2
    positions[id(node)] = (x, y, node["label"])
    if "children" in node:
        total = sum(leaf_count(child) for child in node["children"])
        cursor = x0
        for child in node["children"]:
            width = (x1 - x0) * leaf_count(child) / total
            child_x = cursor + width / 2
            edges.append(((x, y), (child_x, y - 1)))
            layout(child, cursor, cursor + width, y - 1, positions, edges)
            cursor += width


positions, edges = {}, []
layout(semantic_tree, 0, 1, 0, positions, edges)

vocab = [
    "SELECT",
    "name",
    "department",
    "FROM",
    "students",
    "WHERE",
    "age",
    ">",
    "2",
    ";",
    "DROP",
]
allowed_sets = [
    {"SELECT"},
    {"name", "department"},
    {"FROM"},
    {"students"},
    {"WHERE"},
    {"age"},
    {">"},
    {"2"},
    {";"},
]
mask = np.array(
    [[int(token in allowed) for token in vocab] for allowed in allowed_sets]
)