Constrained Decoding › Regex Constraints

import re


class RegexConstraint:
    """
    Constrain generation to match a regular expression.

    Uses a simplified simulation approach for demonstration.
    Real implementations compile to DFA for efficiency.
    """

    def __init__(self, pattern: str):
        self.pattern = pattern
        self.regex = re.compile(f"^{pattern}$")
        self.partial_regex = re.compile(f"^{pattern[:]}")

    def get_valid_tokens(
        self, current: str, vocabulary: list[str]
    ) -> list[str]:
        """
        Find tokens that could extend current string toward a match.

        Args:
            current: String generated so far
            vocabulary: List of possible next tokens

        Returns:
            List of valid next tokens
        """
        valid = []
        for token in vocabulary:
            candidate = current + token
            # Check if candidate could be prefix of a valid match
            if self._could_match(candidate):
                valid.append(token)
        return valid

    def _could_match(self, text: str) -> bool:
        """Check if text could be extended to match the pattern."""
        # Simplified check: see if any extension could work
        # Real implementations use DFA state tracking
        for extension in ["", "a", "0", "@", ".", "-"]:
            test = text + extension * 20  # Pad with potential matches
            if self.regex.match(test):
                return True
        # Also check if text itself is a valid prefix
        # by trying the pattern with .* appended conceptually
        try:
            partial = re.compile(f"^{re.escape(text)}")
            # If we can find any completion that matches, text is valid
            return True  # Simplified - assume prefixes are valid
        except re.error:
            return False

import numpy as np


def generate_with_regex_constraint(
    model_probs: dict[str, float],
    current: str,
    regex_pattern: str,
    vocabulary: list[str],
) -> str:
    """
    Sample next token using both model preferences and regex constraints.

    Args:
        model_probs: Model's probability distribution over tokens
        current: Generated string so far
        regex_pattern: Pattern the final output must match
        vocabulary: List of possible tokens

    Returns:
        Selected token
    """
    constraint = RegexConstraint(regex_pattern)
    valid_tokens = constraint.get_valid_tokens(current, vocabulary)

    if not valid_tokens:
        raise ValueError(f"No valid continuation from '{current}'")

    # Filter model probs to valid tokens only
    filtered_probs = {t: model_probs.get(t, 0) for t in valid_tokens}
    total = sum(filtered_probs.values())

    if total == 0:
        # Uniform over valid tokens if model gives no probability
        return np.random.choice(valid_tokens)

    # Normalize and sample
    normalized = {t: p / total for t, p in filtered_probs.items()}
    tokens, probs = zip(*normalized.items())
    return np.random.choice(tokens, p=probs)