Constrained Decoding

import json

# Examples of common generation failures when asking for structured output
failed_outputs = [
    # Missing closing brace
    '{"name": "Alice", "age": 30',
    # Extra field not in schema
    '{"name": "Bob", "age": 25, "favorite_color": "blue"}',
    # Wrong type (string instead of number)
    '{"name": "Carol", "age": "twenty-five"}',
    # Trailing comma (invalid JSON)
    '{"name": "Dave", "age": 35,}',
    # Mixed format - starts JSON, ends with explanation
    '{"name": "Eve"} Note: I only included the name field since...',
]


def validate_json(text: str) -> tuple[bool, str]:
    """Check if text is valid JSON and return status with message."""
    try:
        json.loads(text)
        return True, "Valid JSON"
    except json.JSONDecodeError as e:
        return False, f"Invalid: {e.msg}"

# 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()

from dataclasses import dataclass
from typing import Any


@dataclass
class JSONSchema:
    """Simplified JSON schema representation."""

    type: str
    properties: dict[str, "JSONSchema"] = None
    required: list[str] = None
    items: "JSONSchema" = None
    enum: list[Any] = None


def schema_to_valid_prefixes(schema: JSONSchema, current: str) -> set[str]:
    """
    Determine valid next characters given schema and current output.

    This is a simplified demonstration. Production implementations
    use finite state machines compiled from the schema.
    """
    current = current.strip()

    if schema.type == "object":
        if not current:
            return {"{"}
        if current == "{":
            if schema.required:
                return {'"'}  # Must start a key
            return {'"', "}"}  # Key or empty object
        # ... more cases for object parsing

    elif schema.type == "string":
        if not current:
            return {'"'}
        if current.startswith('"') and not current.endswith('"'):
            # Inside string - allow any character (simplified)
            return set(
                'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 "'
            )
        if (
            current.startswith('"')
            and current.endswith('"')
            and len(current) > 1
        ):
            return set()  # Complete

    elif schema.type == "number":
        if not current:
            return set("-0123456789")
        if all(c in "-0123456789." for c in current):
            return set("0123456789.")

    elif schema.type == "boolean":
        for val in ["true", "false"]:
            if val.startswith(current):
                if len(current) < len(val):
                    return {val[len(current)]}
                return set()

    return set()

# Example of how Outlines-style generation works conceptually
# (Actual Outlines usage requires installing the library)


def compile_schema_to_fsm(schema: dict) -> dict:
    """
    Compile a JSON schema to a finite state machine.

    Returns a mapping from (state, token) -> next_state.
    This is a simplified conceptual demonstration.
    """
    # In practice, this would build a complete FSM
    # tracking position in the JSON structure
    states = {
        "start": {"valid_tokens": ["{"], "next": "object_open"},
        "object_open": {"valid_tokens": ['"'], "next": "key_start"},
        "key_start": {
            "valid_tokens": list("abcdefghijklmnopqrstuvwxyz"),
            "next": "in_key",
        },
        "in_key": {
            "valid_tokens": list('abcdefghijklmnopqrstuvwxyz"'),
            "next": "key_end_or_continue",
        },
        # ... many more states
    }
    return states


# The FSM approach is much faster than re-parsing at each step
# because state transitions are precomputed

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)

import numpy as np


@dataclass
class Beam:
    """A single beam in constrained beam search."""

    tokens: list[str]
    log_prob: float
    state: Any  # Grammar/FSM state for constraint tracking

    @property
    def text(self) -> str:
        return "".join(self.tokens)


def constrained_beam_search(
    get_token_probs,  # Function returning token probabilities
    constraint_fn,  # Function(text) -> set of valid tokens
    beam_width: int = 3,
    max_length: int = 20,
    vocabulary: list[str] = None,
) -> list[Beam]:
    """
    Beam search with constraint filtering at each step.

    Args:
        get_token_probs: Returns probability distribution for next token
        constraint_fn: Returns valid tokens given current text
        beam_width: Number of beams to maintain
        max_length: Maximum generation length
        vocabulary: Token vocabulary

    Returns:
        Top beams after generation
    """
    if vocabulary is None:
        vocabulary = list("abcdefghijklmnopqrstuvwxyz0123456789-_. ")

    # Initialize with empty beam
    beams = [Beam(tokens=[], log_prob=0.0, state=None)]

    for step in range(max_length):
        all_candidates = []

        for beam in beams:
            # Get valid next tokens
            valid_tokens = constraint_fn(beam.text)
            if not valid_tokens:
                continue  # Dead end

            # Get model probabilities
            probs = get_token_probs(beam.text)

            # Expand beam with valid tokens only
            for token in valid_tokens:
                if token in probs:
                    new_log_prob = beam.log_prob + np.log(probs[token] + 1e-10)
                    new_beam = Beam(
                        tokens=beam.tokens + [token],
                        log_prob=new_log_prob,
                        state=None,
                    )
                    all_candidates.append(new_beam)

        if not all_candidates:
            break

        # Keep top beams by log probability
        all_candidates.sort(key=lambda b: b.log_prob, reverse=True)
        beams = all_candidates[:beam_width]

    return beams

def analyze_constraint_difficulty(
    model_probs: dict[str, float], valid_tokens: set[str]
) -> dict:
    """
    Analyze how much the constraint fights the model's preferences.

    Returns metrics about the constraint's impact on generation.
    """
    valid_prob_mass = sum(model_probs.get(t, 0) for t in valid_tokens)
    invalid_prob_mass = 1.0 - valid_prob_mass

    # Top token analysis
    sorted_tokens = sorted(model_probs.items(), key=lambda x: -x[1])
    top_token, top_prob = sorted_tokens[0]
    top_is_valid = top_token in valid_tokens

    # Find best valid token
    valid_probs = [(t, model_probs.get(t, 0)) for t in valid_tokens]
    valid_probs.sort(key=lambda x: -x[1])
    best_valid = valid_probs[0] if valid_probs else (None, 0)

    return {
        "valid_mass": valid_prob_mass,
        "invalid_mass": invalid_prob_mass,
        "top_token": top_token,
        "top_prob": top_prob,
        "top_is_valid": top_is_valid,
        "best_valid_token": best_valid[0],
        "best_valid_prob": best_valid[1],
        "constraint_difficulty": invalid_prob_mass / (valid_prob_mass + 1e-10),
    }

import numpy as np


def constrained_sample_with_temperature(
    logits: np.ndarray, valid_token_ids: set[int], temperature: float = 1.0
) -> int:
    """
    Sample with both temperature and constraints.

    Args:
        logits: Raw model logits
        valid_token_ids: Set of valid token indices
        temperature: Sampling temperature (>1 = more random, <1 = more deterministic)

    Returns:
        Sampled token ID
    """
    # Apply temperature first
    scaled_logits = logits / temperature

    # Mask invalid tokens
    masked_logits = np.full_like(scaled_logits, float("-inf"))
    for tid in valid_token_ids:
        masked_logits[tid] = scaled_logits[tid]

    # Convert to probabilities
    max_logit = np.max(masked_logits[masked_logits > float("-inf")])
    exp_logits = np.exp(masked_logits - max_logit)
    probs = exp_logits / exp_logits.sum()

    return np.random.choice(len(probs), p=probs)

import numpy as np


def constrained_top_k_sample(
    logits: np.ndarray,
    valid_token_ids: set[int],
    k: int = 5,
    temperature: float = 1.0,
) -> int:
    """
    Top-k sampling within constraint-valid tokens.

    Args:
        logits: Raw model logits
        valid_token_ids: Set of valid token indices
        k: Number of top tokens to consider
        temperature: Sampling temperature

    Returns:
        Sampled token ID
    """
    # First, mask to valid tokens only
    valid_logits = []
    valid_ids_list = []
    for tid in valid_token_ids:
        valid_logits.append(logits[tid])
        valid_ids_list.append(tid)

    valid_logits = np.array(valid_logits)

    # Apply temperature
    valid_logits = valid_logits / temperature

    # Get top-k among valid tokens
    if len(valid_logits) > k:
        top_k_indices = np.argsort(valid_logits)[-k:]
        top_k_logits = valid_logits[top_k_indices]
        top_k_ids = [valid_ids_list[i] for i in top_k_indices]
    else:
        top_k_logits = valid_logits
        top_k_ids = valid_ids_list

    # Convert to probabilities and sample
    exp_logits = np.exp(top_k_logits - np.max(top_k_logits))
    probs = exp_logits / exp_logits.sum()

    selected_idx = np.random.choice(len(top_k_ids), p=probs)
    return top_k_ids[selected_idx]

from typing import Optional


class SimpleSQLConstraint:
    """
    Constraint for simple SQL SELECT statements.

    Supports: SELECT <columns> FROM <table> WHERE <condition>
    """

    KEYWORDS = {"SELECT", "FROM", "WHERE", "AND", "OR"}
    OPERATORS = {"=", ">", "<", ">=", "<=", "!="}

    def __init__(
        self, valid_tables: list[str], valid_columns: dict[str, list[str]]
    ):
        """
        Initialize with schema information.

        Args:
            valid_tables: List of valid table names
            valid_columns: Mapping from table name to list of column names
        """
        self.valid_tables = set(valid_tables)
        self.valid_columns = valid_columns

    def get_valid_tokens(self, current_sql: str) -> set[str]:
        """
        Determine valid next tokens based on SQL parse state.
        """
        tokens = current_sql.upper().split()

        if not tokens:
            return {"SELECT"}

        last_token = tokens[-1]

        if last_token == "SELECT":
            # After SELECT: column names or *
            all_columns = set()
            for cols in self.valid_columns.values():
                all_columns.update(cols)
            return all_columns | {"*"}

        elif (
            last_token in self.valid_columns.get(self._find_table(tokens), [])
            or last_token == "*"
        ):
            # After column: comma for more columns, or FROM
            return {",", "FROM"}

        elif last_token == ",":
            # After comma: another column
            table = self._find_table(tokens)
            if table:
                return set(self.valid_columns.get(table, []))
            return set()

        elif last_token == "FROM":
            # After FROM: table name
            return self.valid_tables

        elif last_token in self.valid_tables:
            # After table: WHERE or end
            return {"WHERE", "<END>"}

        elif last_token == "WHERE":
            # After WHERE: column for condition
            table = self._find_table(tokens)
            return set(self.valid_columns.get(table, []))

        elif last_token in self.OPERATORS:
            # After operator: value (simplified to numbers/quoted strings)
            return {"<VALUE>"}

        return set()

    def _find_table(self, tokens: list[str]) -> Optional[str]:
        """Find the table name in the token list."""
        try:
            from_idx = tokens.index("FROM")
            if from_idx + 1 < len(tokens):
                return tokens[from_idx + 1]
        except ValueError:
            pass
        return None

import numpy as np


def generate_sql_with_constraints(
    constraint: SimpleSQLConstraint,
    model_preferences: dict[str, float],
    max_tokens: int = 20,
) -> str:
    """
    Generate a SQL query using constrained sampling.

    Args:
        constraint: SQL syntax constraint
        model_preferences: Model's token preferences (simulated)
        max_tokens: Maximum query length

    Returns:
        Generated SQL query
    """
    tokens = []

    for _ in range(max_tokens):
        current_sql = " ".join(tokens)
        valid = constraint.get_valid_tokens(current_sql)

        if not valid or "<END>" in valid:
            break

        # Remove <END> and <VALUE> placeholders for actual sampling
        valid = {t for t in valid if not t.startswith("<")}
        if not valid:
            break

        # Sample based on model preferences
        probs = {t: model_preferences.get(t.lower(), 0.1) for t in valid}
        total = sum(probs.values())
        probs = {t: p / total for t, p in probs.items()}

        tokens_list, prob_list = zip(*probs.items())
        selected = np.random.choice(tokens_list, p=prob_list)
        tokens.append(selected)

    return " ".join(tokens)