"""Typed CKY-style parser over locally scored semantic fragments.""" from __future__ import annotations from dataclasses import dataclass, replace import math from strata.modeling.compose.ast import ProgramNode, apply_chain, canonicalize from strata.modeling.compose.fragments import FunctionFragment, LexicalCandidate from strata.modeling.compose.types import SemanticType @dataclass(frozen=True, slots=True) class ChartItem: start: int end: int score: float inside_score: float function: FunctionFragment | None = None program: ProgramNode | None = None derivation: tuple[str, ...] = () def __post_init__(self) -> None: if (self.function is None) == (self.program is None): raise ValueError("chart item must contain exactly one function or value program") @property def output_type(self) -> SemanticType: if self.function is not None: return self.function.output_type assert self.program is not None return self.program.output_type def key(self) -> tuple: if self.program is not None: return ("value", canonicalize(self.program)) assert self.function is not None return ( "function", self.function.argument_side, tuple(operator.value for operator in self.function.operators), ) @dataclass(frozen=True, slots=True) class ParseResult: program: ProgramNode score: float inside_score: float derivation: tuple[str, ...] class TypedChartParser: """Compose only type-compatible fragments; invalid ASTs never enter the chart.""" def __init__(self, *, beam_size: int = 32, combination_score: float = 0.0) -> None: if beam_size <= 0: raise ValueError("beam_size must be positive") self.beam_size = int(beam_size) self.combination_score = float(combination_score) def parse( self, length: int, candidates: tuple[LexicalCandidate, ...], *, required_output: SemanticType | None = None, ) -> ParseResult: if length <= 0: raise ValueError("length must be positive") chart: dict[tuple[int, int], dict[tuple, ChartItem]] = {} for candidate in candidates: if candidate.end > length: raise ValueError("lexical candidate lies outside the chart") item = ChartItem( candidate.start, candidate.end, candidate.score, candidate.score, function=candidate.function, program=candidate.anchor, derivation=(f"lex:{candidate.start}:{candidate.end}",), ) self._insert(chart.setdefault((item.start, item.end), {}), item) for width in range(2, length + 1): for start in range(0, length - width + 1): end = start + width cell = chart.setdefault((start, end), {}) for split in range(start + 1, end): left = chart.get((start, split), {}) right = chart.get((split, end), {}) for left_item in left.values(): for right_item in right.values(): for combined in self._combine(left_item, right_item): self._insert(cell, combined) if len(cell) > self.beam_size: kept = sorted( cell.values(), key=lambda item: (-item.score, repr(item.key())), )[: self.beam_size] chart[(start, end)] = {item.key(): item for item in kept} complete = [item for item in chart.get((0, length), {}).values() if item.program is not None] if required_output is not None: complete = [item for item in complete if item.output_type is required_output] if not complete: raise ValueError("no complete well-typed derivation") best = sorted(complete, key=lambda item: (-item.score, repr(item.key())))[0] assert best.program is not None return ParseResult(best.program, best.score, best.inside_score, best.derivation) def _combine(self, left: ChartItem, right: ChartItem) -> tuple[ChartItem, ...]: score = left.score + right.score + self.combination_score inside = left.inside_score + right.inside_score + self.combination_score derivation = left.derivation + right.derivation + (f"combine:{left.end}",) outputs: list[ChartItem] = [] if left.program is not None and right.function is not None: if right.function.argument_side == "left" and left.output_type is right.function.input_type: outputs.append(ChartItem( left.start, right.end, score, inside, program=apply_chain(left.program, right.function.operators), derivation=derivation, )) if left.function is not None and right.program is not None: if left.function.argument_side == "right" and right.output_type is left.function.input_type: outputs.append(ChartItem( left.start, right.end, score, inside, program=apply_chain(right.program, left.function.operators), derivation=derivation, )) if left.function is not None and right.function is not None: if ( left.function.argument_side == right.function.argument_side == "left" and left.function.output_type is right.function.input_type ): outputs.append(ChartItem( left.start, right.end, score, inside, function=FunctionFragment( left.function.operators + right.function.operators, argument_side="left", ), derivation=derivation, )) if ( left.function.argument_side == right.function.argument_side == "right" and right.function.output_type is left.function.input_type ): outputs.append(ChartItem( left.start, right.end, score, inside, function=FunctionFragment( right.function.operators + left.function.operators, argument_side="right", ), derivation=derivation, )) return tuple(outputs) @staticmethod def _insert(cell: dict[tuple, ChartItem], candidate: ChartItem) -> None: key = candidate.key() previous = cell.get(key) if previous is None: cell[key] = candidate return inside = float(torch_logaddexp(previous.inside_score, candidate.inside_score)) if candidate.score > previous.score: winner = candidate elif candidate.score < previous.score: winner = previous else: winner = min((candidate, previous), key=lambda item: repr(item.derivation)) cell[key] = replace(winner, inside_score=inside) def torch_logaddexp(left: float, right: float) -> float: maximum = max(left, right) if math.isinf(maximum) and maximum < 0: return maximum return maximum + math.log(math.exp(left - maximum) + math.exp(right - maximum)) __all__ = ["ChartItem", "ParseResult", "TypedChartParser"]