transformers-grammar-constraint 0.2.1__tar.gz

transformers_grammar_constraint-0.2.1/PKG-INFO

@@ -0,0 +1,8 @@
+Metadata-Version: 2.3
+Name: transformers-grammar-constraint
+Version: 0.2.1
+Summary: Grammar-constrained LLM token generation via Lark + HuggingFace Transformers
+Requires-Dist: lark>=1.3.1
+Requires-Dist: torch>=2.10.0
+Requires-Dist: transformers>=4.17.0,<6.0.0
+Requires-Python: >=3.12

transformers_grammar_constraint-0.2.1/pyproject.toml

@@ -0,0 +1,40 @@
+[project]
+name = "transformers-grammar-constraint"
+version = "0.2.1"
+description = "Grammar-constrained LLM token generation via Lark + HuggingFace Transformers"
+requires-python = ">=3.12"
+dependencies = [
+    "lark>=1.3.1",
+    "torch>=2.10.0",
+    "transformers>=4.17.0,<6.0.0",
+]
+
+[tool.uv]
+package = true
+exclude-newer = "3d"
+
+
+[build-system]
+requires = ["uv-build>=0.6,<0.12"]
+build-backend = "uv_build"
+
+[tool.pytest.ini_options]
+pythonpath = ["src"]
+addopts = "-m 'not slow'"
+markers = [
+    "slow: requires GPU and network access to download model weights",
+]
+
+[tool.ruff.lint]
+select = ["I"]
+
+[dependency-groups]
+dev = [
+    "pytest>=8.0.0",
+    "hypothesis>=6.0.0",
+    "python-chess>=1.999",
+    "pytest-xdist>=3.0.0",
+    "pre-commit>=4.5.1",
+    "ruff>=0.15.9",
+    "ty>=0.0.29",
+]

transformers_grammar_constraint-0.2.1/src/transformers_grammar_constraint/__init__.py

@@ -0,0 +1,16 @@
+# src/grammar_constrain/__init__.py
+from transformers_grammar_constraint import (
+    compat as _compat,  # noqa: F401 — applies shim at import time
+)
+from transformers_grammar_constraint.grammar import Grammar, LarkGrammar
+from transformers_grammar_constraint.grammars import JsonGrammar, PgnGrammar, SanGrammar
+from transformers_grammar_constraint.processor import GrammarConstrainedLogitsProcessor
+
+__all__ = [
+    "Grammar",
+    "LarkGrammar",
+    "GrammarConstrainedLogitsProcessor",
+    "JsonGrammar",
+    "SanGrammar",
+    "PgnGrammar",
+]

transformers_grammar_constraint-0.2.1/src/transformers_grammar_constraint/compat.py

@@ -0,0 +1,13 @@
+# src/grammar_constrain/compat.py
+"""
+Compatibility shim for transformers 4.x/5.x.
+
+In transformers 4.x, LogitsWarper was a separate base class.
+In 5.x it was removed; warper-style processors now subclass LogitsProcessor directly.
+This module patches the missing name so user code importing LogitsWarper still works.
+"""
+
+import transformers
+
+if not hasattr(transformers, "LogitsWarper"):
+    setattr(transformers, "LogitsWarper", transformers.LogitsProcessor)

transformers_grammar_constraint-0.2.1/src/transformers_grammar_constraint/grammar.py

@@ -0,0 +1,457 @@
+# src/grammar_constrain/grammar.py
+"""
+Grammar base class.
+
+Subclass LarkGrammar and implement `grammar_string` to define a Lark LALR grammar.
+The base class handles Lark parser creation, character-level incremental validation,
+state-based caching, and EOS detection.
+
+`Grammar` is an abstract generic base; `LarkGrammar` is the concrete Lark-backed
+implementation.  DFA-backed grammars (e.g. SanGrammar) subclass `Grammar[_DfaState]`
+directly.
+"""
+
+import logging
+import re
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Generic, TypeVar
+
+from lark import Lark, UnexpectedCharacters, UnexpectedEOF, UnexpectedToken
+from lark.lexer import PatternRE, PatternStr
+
+# Minimal set of characters used as probe suffixes when checking whether a
+# decoded token can be the START of a regex terminal (e.g. is `"` a valid
+# prefix of ESCAPED_STRING?).  We try completing with each suffix in turn;
+# the first successful fullmatch confirms the text is a valid prefix.
+_REGEX_PROBE_SUFFIXES: tuple[str, ...] = (
+    '"',
+    "0",
+    "a",
+    ",",
+    "]",
+    "}",
+    "+",
+    "-",
+    "1",
+    ")",
+)
+
+# Characters probed as single-char continuations in the Earley fast-path
+# pre-filter.  Covers printable ASCII plus common whitespace; non-ASCII
+# first characters fall through and are skipped automatically.
+_EARLEY_PROBE_CHARS: tuple[str, ...] = tuple(
+    chr(i)
+    for i in range(9, 127)  # HT, LF, VT, FF, CR, then space..~
+)
+
+
+@dataclass
+class _TextState:
+    """Lightweight parser state for non-LALR (e.g. Earley) grammars.
+
+    Tracks accumulated text instead of an InteractiveParser object,
+    since parse_interactive() is only available for LALR.
+    """
+
+    text: str = ""
+
+
+@dataclass
+class _LALRState:
+    """Wrapper around an LALR InteractiveParser with explicit text tracking.
+
+    Replaces direct mutation of Lark private attributes (_accumulated_text,
+    _partial_prefix) with first-class dataclass fields, eliminating dependence
+    on Lark's internal API.
+    """
+
+    parser: Any  # InteractiveParser
+    accumulated_text: str = ""
+    partial_prefix: str = ""
+
+
+logger = logging.getLogger(__name__)
+
+_StateT = TypeVar("_StateT")
+
+# Type alias for the Lark-backed state union used by LarkGrammar.
+_LarkState = _LALRState | _TextState
+
+
+class Grammar(ABC, Generic[_StateT]):
+    """Abstract generic base class for grammar-constrained generation.
+
+    Subclass `LarkGrammar` (and implement `grammar_string`) to define a Lark
+    EBNF grammar.  For custom state machines, subclass `Grammar[YourStateType]`
+    directly and implement all abstract methods.
+
+    Example (Lark grammar)::
+
+        class MyGrammar(LarkGrammar):
+            @property
+            def grammar_string(self) -> str:
+                return r'''
+                start: greeting name
+                greeting: "hello" | "hi"
+                name: /[A-Z][a-z]+/
+                '''
+    """
+
+    #: Override in subclasses to use a different Lark parser backend.
+    #: "lalr" is faster but cannot handle ambiguous grammars (e.g. SAN chess).
+    #: "earley" handles ambiguous grammars at the cost of speed.
+    #: "none" disables Lark entirely (used by DFA-backed grammars).
+    lark_parser: str = "lalr"
+
+    @property
+    @abstractmethod
+    def grammar_string(self) -> str:
+        """Return a Lark grammar string (EBNF-like Lark syntax)."""
+        ...
+
+    @abstractmethod
+    def fresh_state(self) -> _StateT:
+        """Return a fresh parser state at the start of the grammar."""
+        ...
+
+    @abstractmethod
+    def advance_state(self, state: _StateT, text: str) -> _StateT:
+        """Return a new state after feeding `text` through the parser.
+
+        Raises ValueError if the text is not a valid continuation.
+        """
+        ...
+
+    @abstractmethod
+    def _state_key(self, state: _StateT) -> tuple:
+        """Return a hashable key representing the current parser state."""
+        ...
+
+    @abstractmethod
+    def _make_state(self, text: str) -> _StateT | None:
+        """Try to create a parser state from `text`.
+
+        Returns a state if text is a valid prefix, None if text is invalid.
+        """
+        ...
+
+    @abstractmethod
+    def _is_accepting(self, state: _StateT) -> bool:
+        """Return True if the parser is in a grammar-complete (accepting) state."""
+        ...
+
+    @abstractmethod
+    def _is_partial_terminal_extension(self, state: _StateT, new_text: str) -> bool:
+        """Return True if new_text extends a partial terminal at the current state."""
+        ...
+
+    @abstractmethod
+    def get_valid_token_ids(
+        self,
+        state: _StateT,
+        tokenizer: Any,
+        *,
+        vocab: dict[str, int] | None = None,
+        decoded_vocab: dict[int, str] | None = None,
+    ) -> frozenset[int]:
+        """Return the set of token IDs that are valid continuations from `state`."""
+        ...
+
+    @abstractmethod
+    def is_valid_complete(self, text: str) -> bool:
+        """Return True if `text` is a complete, valid string for this grammar."""
+        ...
+
+
+class LarkGrammar(Grammar[_LarkState]):
+    """Concrete Grammar backed by Lark LALR/Earley parser.
+
+    Subclass this and implement `grammar_string` to define your grammar.
+
+    Example::
+
+        class MyGrammar(LarkGrammar):
+            @property
+            def grammar_string(self) -> str:
+                return r'''
+                start: greeting name
+                greeting: "hello" | "hi"
+                name: /[A-Z][a-z]+/
+                '''
+    """
+
+    def __init__(self) -> None:
+        if self.lark_parser == "none":
+            return
+        self._parser = Lark(
+            self.grammar_string,
+            parser=self.lark_parser,
+            propagate_positions=False,
+        )
+        # Cache: state_key -> frozenset of valid token ids
+        self._cache: dict[tuple, frozenset[int]] = {}
+        # Precompute per-terminal information used in get_valid_token_ids.
+        # PatternStr terminals: literal string value (for prefix matching).
+        self._terminal_literals: dict[str, str] = {}
+        # PatternRE terminals: compiled regex pattern (for fullmatch + probe checks).
+        self._terminal_patterns: dict[str, re.Pattern[str]] = {}
+        for term in self._parser.terminals:
+            if isinstance(term.pattern, PatternStr):
+                self._terminal_literals[term.name] = term.pattern.value
+            elif isinstance(term.pattern, PatternRE):
+                try:
+                    self._terminal_patterns[term.name] = re.compile(
+                        term.pattern.value, re.DOTALL
+                    )
+                except re.error:
+                    pass
+
+    def fresh_state(self) -> _LarkState:
+        """Return a fresh parser state at the start of the grammar."""
+        if self.lark_parser != "lalr":
+            return _TextState(text="")
+        ip = self._parser.parse_interactive("")
+        ip.exhaust_lexer()
+        return _LALRState(parser=ip)
+
+    def advance_state(self, state: _LarkState, text: str) -> _LarkState:
+        """Return a new state after feeding `text` through the parser.
+
+        Raises ValueError if the text is not a valid continuation.
+        """
+        if isinstance(state, _TextState):
+            new_text = state.text + text
+            new_state = self._make_state(new_text)
+            if new_state is None:
+                raise ValueError(
+                    f"Text {text!r} is not a valid continuation for this grammar state"
+                )
+            return new_state
+        new_text = self._state_text(state) + text
+        new_ip = self._parser.parse_interactive(new_text)
+        try:
+            new_ip.exhaust_lexer()
+        except (UnexpectedCharacters, UnexpectedToken):
+            raise ValueError(
+                f"Text {text!r} is not a valid continuation for this grammar state"
+            )
+        return _LALRState(parser=new_ip, accumulated_text=new_text)
+
+    def _state_key(self, state: _LarkState) -> tuple:
+        """Return a hashable key representing the current parser state."""
+        if isinstance(state, _TextState):
+            return (state.text,)
+        return (tuple(state.parser.parser_state.state_stack), state.partial_prefix)
+
+    def _state_text(self, state: _LarkState) -> str:
+        """Return the accumulated text for a parser state."""
+        if isinstance(state, _TextState):
+            return state.text
+        return state.accumulated_text
+
+    def _make_state(self, text: str) -> _LarkState | None:
+        """Try to create a parser state from `text`.
+
+        Returns a state if text is a valid prefix, None if text is invalid.
+
+        For LALR: uses parse_interactive + exhaust_lexer; UnexpectedEOF means
+        valid but incomplete prefix.
+
+        For Earley: uses parse(); UnexpectedEOF means valid but incomplete prefix.
+        """
+        if self.lark_parser != "lalr":
+            try:
+                self._parser.parse(text)
+                return _TextState(text=text)  # complete valid parse
+            except UnexpectedEOF:
+                return _TextState(text=text)  # valid but incomplete prefix
+            except Exception:
+                return None
+        ip = self._parser.parse_interactive(text)
+        try:
+            ip.exhaust_lexer()
+            return _LALRState(parser=ip, accumulated_text=text)
+        except UnexpectedEOF:
+            return _LALRState(parser=ip, accumulated_text=text)
+        except (UnexpectedCharacters, UnexpectedToken):
+            return None
+
+    def _token_matches_terminal(self, term_name: str, full_candidate: str) -> bool:
+        """Return True if full_candidate is a valid complete or partial match of terminal.
+
+        For PatternStr: checks whether term_value starts with full_candidate (partial
+        prefix) or full_candidate equals term_value (exact match).
+
+        For PatternRE: checks whether full_candidate fully matches the compiled regex
+        OR whether there exists a short probe extension of full_candidate that matches
+        (indicating full_candidate is a valid partial prefix of the terminal).
+        """
+        term_str = self._terminal_literals.get(term_name)
+        if term_str is not None:
+            return term_str.startswith(full_candidate)
+        compiled = self._terminal_patterns.get(term_name)
+        if compiled is not None:
+            if compiled.fullmatch(full_candidate):
+                return True
+            for suffix in _REGEX_PROBE_SUFFIXES:
+                if compiled.fullmatch(full_candidate + suffix):
+                    return True
+        return False
+
+    def _is_partial_terminal_extension(self, state: _LarkState, new_text: str) -> bool:
+        """Return True if new_text extends a partial terminal at the current state.
+
+        Checks whether (partial_prefix + new_text) matches or is a valid prefix of
+        any terminal (PatternStr or PatternRE) expected by the LALR parser at the
+        current state.  This handles the case where a vocabulary token is only part
+        of a multi-character terminal (e.g. the token `"` when the grammar expects
+        the string literal `"age"` or the regex terminal ESCAPED_STRING).
+        """
+        if isinstance(state, _TextState):
+            return False
+        full_suffix = state.partial_prefix + new_text
+        if not full_suffix:
+            return False
+        try:
+            expected = state.parser.accepts()
+        except Exception:
+            return False
+        return any(self._token_matches_terminal(t, full_suffix) for t in expected)
+
+    def _is_accepting(self, state: _LarkState) -> bool:
+        """Return True if the parser is in a grammar-complete (accepting) state."""
+        if isinstance(state, _TextState):
+            try:
+                self._parser.parse(state.text)
+                return True
+            except Exception:
+                return False
+        try:
+            accepts = state.parser.accepts()
+            return "$END" in accepts
+        except Exception:
+            return False
+
+    def get_valid_token_ids(
+        self,
+        state: _LarkState,
+        tokenizer: Any,
+        *,
+        vocab: dict[str, int] | None = None,
+        decoded_vocab: dict[int, str] | None = None,
+    ) -> frozenset[int]:
+        """Return the set of token IDs that are valid continuations from `state`.
+
+        Includes EOS token if the grammar is currently in an accepting state.
+        Uses state-based caching: identical LALR states reuse prior results.
+
+        Args:
+            state: Current InteractiveParser state (from fresh_state or advance_state).
+            tokenizer: HuggingFace tokenizer with get_vocab() and decode() methods.
+
+        Returns:
+            frozenset of valid token IDs. Empty frozenset signals a dead end.
+
+        Note:
+            For LALR grammars, token-to-terminal matching uses a probe-suffix heuristic to
+            detect valid regex-terminal prefixes (see `_REGEX_PROBE_SUFFIXES`). The probe
+            set is finite and may miss unusual terminals — tokens that are valid prefixes of
+            an uncovered regex pattern will be incorrectly excluded.
+        """
+        state_key = self._state_key(state)
+
+        if state_key in self._cache:
+            return self._cache[state_key]
+
+        _vocab: dict[str, int] = vocab if vocab is not None else tokenizer.get_vocab()
+        valid_ids: set[int] = set()
+
+        if isinstance(state, _TextState):
+            # Earley / non-LALR: fall back to full re-parse for each candidate.
+            current_text = self._state_text(state)
+
+            # Fast pre-filter: probe ASCII characters one at a time to find
+            # which first characters lead to valid continuations.  Tokens
+            # whose decoded form starts with an invalid character are skipped
+            # without a full re-parse, cutting O(vocab) Earley calls to
+            # O(128 + |filtered_vocab|) — critical for large vocabularies.
+            valid_first_chars: frozenset[str] = frozenset(
+                c
+                for c in _EARLEY_PROBE_CHARS
+                if self._make_state(current_text + c) is not None
+            )
+
+            for token_str, token_id in _vocab.items():
+                if not token_str:
+                    continue
+                decoded = (
+                    decoded_vocab[token_id]
+                    if decoded_vocab is not None
+                    else tokenizer.decode([token_id])
+                )
+                if not decoded:
+                    continue
+                # Skip tokens whose first character cannot start a valid
+                # continuation.  Non-ASCII first characters are also filtered
+                # because they will not be present in valid_first_chars.
+                if decoded[0] not in valid_first_chars:
+                    continue
+                if self._make_state(current_text + decoded) is not None:
+                    valid_ids.add(token_id)
+        else:
+            # LALR: use state.accepts() to determine the expected terminals, then
+            # check each vocabulary token against those terminals directly.  This
+            # avoids the false-positive that arises when _make_state() re-parses
+            # the whole accumulated text and the greedy lexer produces a different
+            # tokenisation (e.g. "00" instead of "0" + "0") that happens to be
+            # valid even though the incremental parser has already committed past
+            # the token boundary.
+            partial = state.partial_prefix
+            try:
+                expected_terms = frozenset(state.parser.accepts())
+            except Exception:
+                expected_terms = frozenset()
+
+            for token_str, token_id in _vocab.items():
+                if not token_str:
+                    continue
+                decoded = (
+                    decoded_vocab[token_id]
+                    if decoded_vocab is not None
+                    else tokenizer.decode([token_id])
+                )
+                if not decoded:
+                    continue
+                full_candidate = partial + decoded
+                if any(
+                    self._token_matches_terminal(t, full_candidate)
+                    for t in expected_terms
+                ):
+                    valid_ids.add(token_id)
+
+        # Include EOS if grammar is in accepting state
+        if self._is_accepting(state) and hasattr(tokenizer, "eos_token_id"):
+            eos_id = tokenizer.eos_token_id
+            if eos_id is not None:
+                valid_ids.add(eos_id)
+
+        if not valid_ids:
+            logger.warning(
+                "Grammar reached a dead end: no valid tokens. "
+                "Allowing all tokens to prevent hanging generation."
+            )
+            result = frozenset(_vocab.values())
+        else:
+            result = frozenset(valid_ids)
+
+        self._cache[state_key] = result
+        return result
+
+    def is_valid_complete(self, text: str) -> bool:
+        """Return True if `text` is a complete, valid string for this grammar."""
+        try:
+            self._parser.parse(text)
+            return True
+        except Exception:
+            return False

transformers_grammar_constraint-0.2.1/src/transformers_grammar_constraint/grammars/__init__.py

@@ -0,0 +1,6 @@
+# src/grammar_constrain/grammars/__init__.py
+from transformers_grammar_constraint.grammars.json_grammar import JsonGrammar
+from transformers_grammar_constraint.grammars.pgn_grammar import PgnGrammar
+from transformers_grammar_constraint.grammars.san_grammar import SanGrammar
+
+__all__ = ["JsonGrammar", "SanGrammar", "PgnGrammar"]