constrainers.py

from dataclasses import dataclass
from abc import ABC, abstractmethod
from transformers_cfg.token_grammar_recognizer import AbsTokenRecognizer, check_token_acceptance_in_trie
from transformers_cfg.recognizer import AcceptState
from copy import deepcopy
from functools import lru_cache
import numpy as np

class Constrainer(ABC):
  @abstractmethod
  def get_blocklist(self, tokens: list[int]) -> list[int]:
    ...
  def ask_token(self, tokens: list[int], token: int) -> bool:
    return token not in self.get_blocklist(tokens)

@dataclass(frozen=True)
class TokenRecognizerConstrainer(Constrainer):
  recognizer: AbsTokenRecognizer
  max_gen_length: int

  @lru_cache
  def get_length(self):
    return len(self.recognizer.tokenizer)

  def get_accept_state(self, tokens: list[int]) -> AcceptState:
    initial_state = deepcopy(self.recognizer.string_recognizer.get_initial_parsing_state())
    string = self.recognizer.tokenizer.decode(tokens)
    return self.recognizer.string_recognizer._update_state_with_string(string, initial_state)

  @lru_cache(maxsize=1024)
  def get_token_acceptance_array_for_stack(self, stack: tuple[int]):
    accepts = np.zeros(self.get_length(), dtype=bool)
    token_acceptance = check_token_acceptance_in_trie(
        self.recognizer.byte_trie.root,
        [list(stack)],
        self.recognizer.string_recognizer,
        self.recognizer.eos_token_id,
        accepts,
    )
    assert not token_acceptance[self.recognizer.eos_token_id]
    if not token_acceptance.any():
      token_acceptance[self.recognizer.eos_token_id] = True
    return token_acceptance

  def get_blocklist(self, tokens: list[int]) -> list[int]:
    accept_state = self.get_accept_state(tuple(tokens))
    if not accept_state.stacks:
      return [i for i in range(self.get_length()) if i != self.recognizer.eos_token_id]
    acceptance_matrix = [
      self.get_token_acceptance_array_for_stack(
        tuple(stack)
      )
      for stack in accept_state.stacks
    ]
    blocklist = (~np.stack(acceptance_matrix).any(axis=0)).nonzero()[0].tolist()
    return blocklist