Source code for lambeq.bobcat.parser

# Copyright 2021-2023 Cambridge Quantum Computing Ltd.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from import Iterable, Mapping
from dataclasses import dataclass, field
import math
from typing import overload, Tuple

from lambeq.bobcat.grammar import Grammar
from lambeq.bobcat.lexicon import Atom, Category
from lambeq.bobcat.lexicon import CATEGORIES
from lambeq.bobcat.rules import Rules
from lambeq.bobcat.tree import Dependency, Lexical, ParseTree

SpanT = Tuple[int, int]

NEGATIVE_INFINITY = -float('inf')

[docs]@dataclass class Supertag: """A string category, annotated with its log probability.""" category: str probability: float
[docs]@dataclass class Sentence: """An input sentence. Attributes ---------- words : list of str The tokens in the sentence. input_supertags : list of list of Supertag A list of supertags for each word. span_scores : dict of tuple of int and int to dict of int to float Mapping of a span to a dict of category (indices) mapped to their log probability. """ words: list[str] input_supertags: list[list[Supertag]] span_scores: dict[SpanT, dict[int, float]] def __post_init__(self) -> None: if len(self.words) != len(self.input_supertags): raise ValueError( '`words` must be the same length as `input_supertags`') def __len__(self) -> int: return len(self.words)
@dataclass class Cell: """A cell in the chart. The cell maintains a list of trees in sorted order, up to the beam size (though may be larger if there are ties at the bottom), with the further restriction that only one tree is allowed per category. """ beam_size: int trees: list[ParseTree] = field(default_factory=list) trees_map: dict[Category, ParseTree] = field(default_factory=dict) min_score: float = NEGATIVE_INFINITY def find(self, score: float) -> int: """Find the index where a tree with the given score can go.""" trees = self.trees lo = 0 hi = len(trees) while lo < hi: mid = (lo + hi) // 2 cmp = trees[mid].score if score == cmp: return mid elif score > cmp: hi = mid else: lo = mid + 1 return lo def add(self, to_add: Iterable[ParseTree]) -> int: """Add the trees to the cell. For each tree that is to be added, it is checked against the existing trees to determine whether it should be added, and if so, is added using a binary search; then, the beam is applied. """ to_add = sorted(to_add, key=lambda tree: -tree.score) trees = self.trees trees_map = self.trees_map added = 0 b = self.beam_size for tree in to_add: score = tree.score if len(trees) >= b and score < trees[-1].score: break # Check whether there exists a tree with the same category. # If there does, and it has a lower score, then remove the # old tree before inserting the new tree. # If the score is higher, then do nothing. insert: bool try: old_tree = trees_map[] except KeyError: insert = True else: old_score = old_tree.score insert = score > old_score if insert: old_index = self.find(old_score) deleted = False for i in range(old_index, len(trees)): if trees[i] is old_tree: del trees[i] deleted = True break elif trees[i].score != old_score: break if not deleted: for i in reversed(range(old_index)): if trees[i] is old_tree: del trees[i] break if insert: trees.insert(self.find(score), tree) trees_map[] = tree added += 1 try: cutoff = self.min_score = trees[b - 1].score if trees[b].score < cutoff: added -= len(trees) - b for tree in trees[b:]: del trees_map[] del trees[b:] except IndexError: pass return added @dataclass class Chart: """The parse chart, containing a mapping from span to cell. A span (i, j) represents the phrase from the ith word to the jth word (inclusive), indexed from 0. """ beam_size: int chart: dict[SpanT, Cell] = field(default_factory=dict) parse_tree_count: int = 0 def __getitem__(self, index: SpanT) -> list[ParseTree]: return self.chart[index].trees def min_score(self, start: int, end: int) -> float: """Get the lowest score needed to add a tree to the given cell.""" try: return self.chart[start, end].min_score except KeyError: return NEGATIVE_INFINITY def add(self, start: int, end: int, to_add: Iterable[ParseTree]) -> None: """Add parse trees to the cell in the chart.""" if not to_add: return try: cell = self.chart[start, end] except KeyError: cell = self.chart[start, end] = Cell(self.beam_size) self.parse_tree_count += cell.add(to_add) @dataclass class ParseResult: """The result of a parse. This acts as a list of the most probable parse trees, in order, i.e. use `parse_result[0]` to access the most probable parse tree. Parameters ---------- chart : Chart The parse chart. Attributes ---------- words : list[str] The words in the sentence. root : list[str] The most probable parse trees, in order. """ chart: Chart words: list[str] = field(init=False) root: list[ParseTree] = field(init=False) _output_tags: list[Category] = field(default_factory=list, init=False) def __post_init__(self) -> None: self.words = [] while True: try: tree = self.chart[len(self.words), len(self.words)][0] except KeyError: break else: while tree.left: tree = tree.left self.words.append(tree.word) try: self.root = self.chart[0, len(self.words) - 1] except KeyError: self.root = [] def __bool__(self) -> bool: return len(self.root) != 0 def __len__(self) -> int: return len(self.root) @overload def __getitem__(self, index: int) -> ParseTree: ... @overload def __getitem__(self, index: slice) -> list[ParseTree]: ... def __getitem__(self, index: int | slice) -> ParseTree | list[ParseTree]: return self.root[index] def deps( self, tree: ParseTree | None = None ) -> tuple[list[Dependency], list[str]]: # pragma: no cover """Get the dependencies and output tags of the parse. If `tree` is not specified, then this looks for the best scoring tree at the root of the parse; if there is none, then it amalgamates results from the best-scoring trees in the chart. """ if tree is None: try: tree = self.root[0] except IndexError: return self._skim_deps() deps = tree.filled_deps.copy() tags = [] if tree.left: for child in (tree.left, tree.right): if child: child_deps, child_tags = self.deps(child) deps += child_deps tags += child_tags else: tags.append(str('[X]', '')) return deps, tags def _skim_deps( self, start: int = 0, end: int | None = None ) -> tuple[list[Dependency], list[str]]: # pragma: no cover if end is None: end = len(self.words) - 1 if start > end: return [], [] result_start = result_end = max_tree = None for span_length in reversed(range(end + 1 - start)): max_score = NEGATIVE_INFINITY for i in range(start, end + 1 - span_length): try: cell = self.chart[i, i + span_length] except KeyError: pass else: tree = cell[0] if tree.score > max_score: max_score = tree.score max_tree = tree result_start = i result_end = i + span_length if max_tree: break left_deps, left_tags = self._skim_deps(start, result_start - 1) tree_deps, tree_tags = self.deps(max_tree) right_deps, right_tags = self._skim_deps(result_end + 1, end) return (left_deps + tree_deps + right_deps, left_tags + tree_tags + right_tags)
[docs]class ChartParser:
[docs] def __init__(self, grammar: Grammar, cats: Iterable[str], root_cats: Iterable[str] | None, eisner_normal_form: bool, max_parse_trees: int, beam_size: int, input_tag_score_weight: float, missing_cat_score: float, missing_span_score: float) -> None: self.max_parse_trees = max_parse_trees self.categories = {} for plain_cat, markedup_cat in grammar.categories.items(): self.categories[plain_cat] = Category.parse(markedup_cat) self.rules = Rules(eisner_normal_form, grammar, self.categories) self.input_tag_score_weight = input_tag_score_weight self.beam_size = beam_size try: self.missing_cat_score = math.log(missing_cat_score) except ValueError: self.missing_cat_score = NEGATIVE_INFINITY try: self.missing_span_score = math.log(missing_span_score) except ValueError: self.missing_span_score = NEGATIVE_INFINITY CONJ_TAG = '[conj]' self.result_cats: dict[tuple[str, tuple[Category, ...]], int] = {} cat_id = 0 for cat_str in cats: chain = cat_str.split('::') res_cats: tuple[Category, ...] if len(chain) == 1 and chain[0].endswith(CONJ_TAG): base_cat = chain[0][:-len(CONJ_TAG)] if '/' in base_cat or '\\' in base_cat: base_cat = f'({base_cat})' cat_modified = fr'({base_cat}\{base_cat})' res_cats = (Category.parse(cat_modified),) label = 'conj' else: res_cats = tuple(map(Category.parse, chain)) label = 'unary' if len(chain) > 1 else 'binary' self.result_cats[label, res_cats] = cat_id cat_id += 1 self.set_root_cats(root_cats)
[docs] def set_root_cats(self, root_cats: Iterable[Category | str] | None) -> None: if root_cats is None: self.root_cats = None else: try: self.root_cats = [(cat if isinstance(cat, Category) else CATEGORIES[cat, 0]) for cat in root_cats] except KeyError as e: raise ValueError('Grammar does not contain root category: ' f'{repr(e.args[0])}') from e
[docs] def filter_root(self, trees: list[ParseTree]) -> list[ParseTree]: if self.root_cats is None: return trees else: results = [] for tree in trees: for cat in self.root_cats: if cat.matches( results.append(tree) break return results
[docs] def __call__(self, sentence: Sentence) -> ParseResult: """Parse a sentence.""" chart = Chart(self.beam_size) for i, (word, supertags) in enumerate(zip(sentence.words, sentence.input_supertags)): results = [] for supertag in supertags: tree = Lexical(self.categories[supertag.category], word, i + 1) tree.score = self.input_tag_score_weight * supertag.probability results.append(tree) try: span_scores = sentence.span_scores[i, i] except KeyError: pass else: if len(sentence) > 1: results += self.rules.type_change(results) results += self.rules.type_raise(results) for tree in results: if tree.left: self.calc_score_unary(tree, span_scores) # filter root cats if len(sentence) == 1: results = self.filter_root(results) chart.add(i, i, results) for span_length in range(1, len(sentence)): for end in range(span_length, len(sentence)): if chart.parse_tree_count > self.max_parse_trees: break start = end - span_length try: span_scores = sentence.span_scores[start, end] except KeyError: continue max_span_score = max((self.missing_cat_score, self.missing_span_score, *span_scores.values())) for split in range(start + 1, end + 1): try: left_trees = chart[start, split - 1] right_trees = chart[split, end] except KeyError: continue for left in left_trees: for right in right_trees: max_score = (left.score + right.score + max_span_score) if max_score < chart.min_score(start, end): break results = self.rules.combine(left, right) if results and len(sentence) > span_length + 1: results += self.rules.type_change(results) results += self.rules.type_raise(results) # filter root cats if span_length == len(sentence) - 1: results = self.filter_root(results) for tree in results: if tree.right: self.calc_score_binary(tree, span_scores) else: self.calc_score_unary(tree, span_scores) chart.add(start, end, results) return ParseResult(chart)
[docs] def calc_score_unary(self, tree: ParseTree, span_scores: Mapping[int, float]) -> None: """Calculate the score for a unary tree (chain).""" left = tree.left res_cat: tuple[str, tuple[Category, ...]] if left.right is None and left.left is not None: base = left.left res_cat = ('unary', (,, else: base = left res_cat = ('unary', (, if base.right is not None: tree.score = base.left.score + base.right.score else: tree.score = base.score cat_id = self.result_cats.get(res_cat) tree.score += self.get_span_score(span_scores, cat_id)
[docs] def calc_score_binary(self, tree: ParseTree, span_scores: Mapping[int, float]) -> None: """Calculate the score for a binary tree.""" if tree.coordinated: cat_id = self.result_cats.get(('conj', (,))) else: cat = try: cat_id = self.result_cats['binary', (,)] except KeyError: if cat.atom == Atom.NP: cat_no_nb = Category(cat.atom) cat_id = self.result_cats.get(('binary', (cat_no_nb,))) else: cat_id = None tree.score = (tree.left.score + tree.right.score + self.get_span_score(span_scores, cat_id))
[docs] def get_span_score(self, span_scores: Mapping[int, float], cat_id: int | None) -> float: """Get the score in a span for a category (chain) ID.""" if cat_id is None: return self.missing_cat_score try: return span_scores[cat_id] except KeyError: return self.missing_span_score