source/src/strata/modeling/compose/chart_parser.py

"""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"]