masking.py

"""Shared grammar-legality computation + a per-state cache.

The set of grammar-legal next tokens is a pure function of the grammar state, so
we cache the boolean legal mask by state. This is what makes the crossing search
affordable: its many short rollouts all start from the same handful of
line-start states and reuse one (expensive) full-vocab scan.
"""

from __future__ import annotations

import numpy as np
import torch


class LegalCache:
    def __init__(self, grammar, token_text, eos_token_ids=()):
        self.grammar = grammar
        self.token_text = token_text
        self.eos_token_ids = [int(x) for x in eos_token_ids]
        # Special tokens decode to '' and are always illegal — never probe them.
        self._scan_ids = [i for i, t in enumerate(token_text) if t]
        self._legal_cache = {}      # state-key -> np.bool_ array
        self._illegal_cache = {}    # state-key -> torch.BoolTensor

    @staticmethod
    def _key(state):
        return state if isinstance(state, int) else tuple(state)

    def legal_np(self, state):
        key = self._key(state)
        cached = self._legal_cache.get(key)
        if cached is not None:
            return cached
        advance = self.grammar.advance
        token_text = self.token_text
        at_accept = self.grammar.accepts(state)
        legal = np.zeros(len(token_text), dtype=bool)
        for i in self._scan_ids:
            if advance(state, token_text[i]) != -1:
                legal[i] = True
        for eid in self.eos_token_ids:
            legal[eid] = at_accept
        self._legal_cache[key] = legal
        return legal

    def illegal_tensor(self, state):
        key = self._key(state)
        cached = self._illegal_cache.get(key)
        if cached is not None:
            return cached
        t = torch.from_numpy(~self.legal_np(state))
        self._illegal_cache[key] = t
        return t