Constrained Decoding › Grammar-Guided Generation

# A simple grammar for arithmetic expressions
# E -> E + T | E - T | T
# T -> T * F | T / F | F
# F -> ( E ) | number

# For constrained decoding, we need to track parser state
class SimpleArithmeticParser:
    """
    Track valid next tokens for simple arithmetic expressions.

    This is a simplified example - real implementations use
    more sophisticated parsing techniques.
    """

    def __init__(self):
        self.stack = ["E"]  # Start with expression
        self.depth = 0  # Parenthesis depth

    def valid_next_tokens(self, current_text: str) -> set[str]:
        """Determine valid next tokens based on current parse state."""
        # Simplified logic for demonstration
        if not current_text:
            return {"(", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9"}

        last_char = current_text[-1]

        if last_char.isdigit():
            # After a number: operators, closing paren, or more digits
            valid = {
                "+",
                "-",
                "*",
                "/",
                "0",
                "1",
                "2",
                "3",
                "4",
                "5",
                "6",
                "7",
                "8",
                "9",
            }
            if self.depth > 0:
                valid.add(")")
            return valid

        elif last_char in "+-*/":
            # After operator: number or opening paren
            return {"(", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9"}

        elif last_char == "(":
            self.depth += 1
            return {"(", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9"}

        elif last_char == ")":
            self.depth -= 1
            valid = {"+", "-", "*", "/"}
            if self.depth > 0:
                valid.add(")")
            return valid

        return set()

import numpy as np


class GrammarConstrainedSampler:
    """
    Apply grammar constraints during text generation.

    This simplified implementation assumes character-level tokenization.
    Real implementations must handle subword tokenization carefully.
    """

    def __init__(self, grammar_valid_fn):
        """
        Initialize with a function that returns valid next tokens.

        Args:
            grammar_valid_fn: Function(current_text) -> set of valid tokens
        """
        self.grammar_valid_fn = grammar_valid_fn

    def constrained_sample(
        self,
        logits: np.ndarray,
        current_text: str,
        token_to_char: dict[int, str],
    ) -> int:
        """
        Sample from logits with grammar constraints.

        Args:
            logits: Raw model logits over vocabulary
            current_text: Text generated so far
            token_to_char: Mapping from token IDs to characters

        Returns:
            Selected token ID
        """
        # Get valid tokens from grammar
        valid_chars = self.grammar_valid_fn(current_text)

        # Create mask for valid token IDs
        valid_token_ids = {
            tid for tid, char in token_to_char.items() if char in valid_chars
        }

        # Apply mask: set invalid logits to -infinity
        masked_logits = logits.copy()
        for tid in range(len(logits)):
            if tid not in valid_token_ids:
                masked_logits[tid] = float("-inf")

        # Convert to probabilities and sample
        probs = self._softmax(masked_logits)
        return np.random.choice(len(probs), p=probs)

    def _softmax(self, x: np.ndarray) -> np.ndarray:
        """Compute softmax probabilities."""
        exp_x = np.exp(x - np.max(x))
        return exp_x / exp_x.sum()