transformers_gad/token_grammar_recognizer.py

import copy
import logging
from abc import ABC
from functools import lru_cache
from typing import List

import torch

from transformers_gad.recognizer import StringRecognizer, AcceptState
from transformers_gad.parser import parse_ebnf
from transformers_gad.trie import ByteTrie
from transformers_gad.utf8_utils import PartialUTF8
from .vocab_struct import LEAF, TokenTrie
from transformers_gad.mapping import get_mapping

logger = logging.getLogger(__name__)


class AbsTokenRecognizer(ABC):
    def __init__(self, grammar_str, tokenizer, start_rule_name="root", unicode=False):
        parsed_grammar = parse_ebnf(grammar_str)
        grammar_encoding = parsed_grammar.grammar_encoding
        self.start_rule_id = parsed_grammar.symbol_table.get(start_rule_name)
        self.byte_encoding = unicode

        if unicode and not tokenizer.__class__.__name__.lower().startswith(
            "gpt2"
        ):  # gpt2tokenizer or gpt2tokenizerfast
            raise ValueError(
                "Constrained decoding with unicode is only supported for GPT2 model. Support for other models is coming soon."
                "Or you can use the constraints with only ascii characters."
            )

        self.eos_token_id = tokenizer.eos_token_id
        self.token_trie = TokenTrie(tokenizer)
        self.tokenizer = tokenizer
        self.string_recognizer = StringRecognizer(grammar_encoding, self.start_rule_id)
        self.unicode_trie = ByteTrie.from_tokenizer(tokenizer, unicode=unicode)
        self.mapping = get_mapping(tokenizer, unicode=unicode)
        assert len(self.mapping) == len(
            self.token_trie
        ), f"{len(self.mapping)}, {len(self.token_trie)}"

    def _consume_token_id(
        self, token_id: int, accept_state: AcceptState
    ) -> AcceptState:
        if self.string_recognizer._must_stop(accept_state.stacks):
            if token_id == self.eos_token_id:
                return self.string_recognizer.get_termination_accept_state()
            else:
                raise ValueError(
                    f"All stacks are empty, so the only token accepted is EOS({self.eos_token_id}), but got {token_id}"
                )
        if token_id == self.eos_token_id:
            if self.string_recognizer._can_stop(accept_state.stacks):
                # if at least one of the stack is empty, we can stop
                # we clear all the stacks, meaning that we don't accept any token after EOS
                return self.string_recognizer.get_termination_accept_state()
            else:
                raise ValueError(
                    f"At least one of the stack should be empty when EOS is reached. However, "
                    f"the stacks are {accept_state.stacks}"
                )

        bytes_or_codepoints = self.mapping.map(token_id)
        accept_state = self.string_recognizer._consume_bytes(
            bytes_or_codepoints, accept_state
        )
        return accept_state

    def probe_token_id(self, token_id: int, accept_state: AcceptState) -> bool:
        stacks = accept_state.stacks
        if self.string_recognizer._must_stop(stacks):
            if token_id == self.eos_token_id:
                return True
            else:
                return False
        if token_id == self.eos_token_id:
            if self.string_recognizer._can_stop(stacks):
                # if at least one of the stack is empty, we can stop
                # we clear all the stacks, meaning that we don't accept any token after EOS
                return True
            else:
                return False
        # for code_point in self.mapping.map(token_id):
        #     stacks = self.grammar._consume_char_code_point(code_point, stacks)
        bytes_or_codepoints = self.mapping.map(token_id, verbose=False)
        new_acc_state = self.string_recognizer._consume_bytes(
            bytes_or_codepoints, accept_state, verbose=False
        )
        return len(new_acc_state.stacks) > 0

    def advance_token_ids(self, *args, **kwargs):
        """Process a list of tokens according to the grammar rules."""
        raise NotImplementedError

    def batch_filter_vocab(self, batch_accept_states, device) -> torch.Tensor:
        batch_acceptance = []
        for accept_state in batch_accept_states:
            batch_acceptance.append(self.filter_vocab(accept_state, device))
        return torch.stack(batch_acceptance)

    def filter_vocab(self, accept_state, device) -> torch.Tensor:
        if not accept_state.stacks:  # Check if stacks is empty
            # Handle the empty case: for example, return a tensor of False
            # The size of the tensor should match the size of your vocabulary
            vocab_size = len(self.mapping)
            logger.debug(f"Empty stack, sum of acceptance: {0}")
            # size of the vocab
            accepts = [False] * vocab_size
            accepts[self.eos_token_id] = True
            return torch.tensor(accepts, dtype=torch.bool, device=device)

        acceptance = self.get_token_acceptance(accept_state, device)

        return acceptance

    def get_token_acceptance(self, accept_state, device) -> torch.Tensor:
        acceptance_matrix = torch.cat(
            [
                self.get_token_acceptance_array_for_stack(
                    tuple(stack), accept_state.partial_utf8, device
                )
                for stack in accept_state.stacks
            ]
        )
        # Merge stacks: any True => True
        acceptance = acceptance_matrix.reshape(len(accept_state.stacks), -1).any(dim=0)
        return acceptance

    @lru_cache(maxsize=32768)
    def get_token_acceptance_array_for_stack(self, stack, partial_utf8, device):
        # stack = list(stack)  # needs to come in as a tuple for lru_cache
        assert isinstance(stack, tuple)
        stack = list(stack)

        if self.byte_encoding:

            accept_f = lambda x: self.string_recognizer._probe_bytes(
                x, [stack], partial_utf8=partial_utf8
            )
            token_acceptance = self.unicode_trie.get_token_acceptance(
                accept=accept_f, accept_eos=False, eos_token_id=self.eos_token_id
            )
        else:
            accepts = [False] * len(self.mapping)
            token_acceptance = check_token_acceptance_in_trie(
                self.token_trie.trie,
                [stack],
                self.string_recognizer,
                self.eos_token_id,
                accepts,
            )
        x = torch.tensor(token_acceptance, dtype=torch.bool, device=device)
        x_eos = self.validate_and_set_eos_acceptance(x)
        return x_eos

    def validate_and_set_eos_acceptance(self, acceptance: torch.Tensor) -> torch.Tensor:
        if torch.any(acceptance) == 0:
            acceptance[self.eos_token_id] = True
        else:
            if acceptance[self.eos_token_id]:
                raise ValueError()
            acceptance[self.eos_token_id] = False
        return acceptance


class IncrementalTokenRecognizer(AbsTokenRecognizer):
    def __init__(self, grammar_str, start_rule_name, tokenizer, unicode=False):
        super().__init__(grammar_str, tokenizer, start_rule_name, unicode)
        self.last_size = None
        self.is_incremental = True

        # if self.last_size is not set (which would be the case when processing the first token).
        # In this case, do nothing.

    def advance_token_ids(self, input_ids, batch_accept_states, parse_start_index=None):

        if self.last_size is None:
            prefix_to_parse = [
                single_input_ids[parse_start_index:]
                if parse_start_index is not None
                else []
                for single_input_ids in input_ids
            ]

            # self.grammar_acceptor.accept_token_ids(prefix_to_parse, self.stacks)
            batch_accept_states = [
                self._consume_token_ids(prefix, accept_state)
                for prefix, accept_state in zip(prefix_to_parse, batch_accept_states)
            ]
            #  if the length of the current input IDs (input_ids[0]) is exactly one more than self.last_size.
            #  This is expected in a scenario where inputs are processed incrementally, one token at a time.
        elif len(input_ids[0]) == self.last_size + 1:
            batch_accept_states = [
                self._consume_token_id(single_input_ids[-1], accept_state)
                for single_input_ids, accept_state in zip(
                    input_ids, batch_accept_states
                )
            ]
            #  ensure that the input size is consistent with the expected incremental processing
            #  (i.e., one token at a time).
        else:
            # here we check if the input_ids are one token longer than the last time we processed
            # but we don't check if input_ids are actually valid.
            # Imagine a scenario where we generate 10 tokens, then we replace the 10 generated tokens with 10 new tokens.
            # In this case, the input_ids will be consistent with the last_size, but the input_ids are not valid.
            # However, should we really check if the input_ids are valid here?
            # If we do, then we need to reparse the whole input_ids at each call, which is not efficient.
            # Maybe we should just trust the user to provide valid input_ids?
            # The conclusion is that, we assume the input_ids are valid, and our generation will be correct.
            # If the input_ids are not valid, then the generation result will be wrong and we don't take responsibility for that.
            raise RuntimeError(
                "Input ID's length is inconsistent with the current state of "
                "the GrammarConstrainedLogitsProcessor. If you want to process "
                "another input sequence, please instantiate a new "
                "GrammarConstrainedLogitsProcessor " 
                "or call reset_parser method of GrammarAlignedOracleLogitsProcessor"
            )
        self.last_size = len(input_ids[0])

        return batch_accept_states

    def _consume_token_ids(
        self, token_ids: List[int], accept_state: AcceptState = None, as_string=True
    ):
        if accept_state is None:
            accept_state = self.string_recognizer.get_initial_accept_state()
        if as_string:
            string = self.tokenizer.decode(token_ids)
            accept_state = self.string_recognizer._consume_string(string, accept_state)
        else:
            for i, token_id in enumerate(token_ids):
                accept_state = self._consume_token_id(token_id, accept_state)
                if len(accept_state.stacks) > 0:
                    cur_token_ids = token_ids[: i + 1]
                    logging.debug(f"{cur_token_ids} is accepted")
                    decoded_string = self.tokenizer.decode(cur_token_ids)
                    logging.debug(f"The decoded string is {decoded_string}")
        return accept_state

    def reset(self):
        self.last_size = None

def check_token_acceptance_in_trie(trie, stacks, grammar, eos_token_id, accepts):

    for byte, next_trie in trie.items():
        if byte == LEAF:
            token_id = next_trie
            if token_id != eos_token_id:
                # if the stacks is not empty, it means we can still continue to parse
                # so we should accept the token
                accepts[token_id] = bool(stacks)
            continue

        new_stacks = []
        for stk in stacks:
            if not stk:
                continue

            next_element_offset = stk[-1]
            num_chars = grammar.grammar_encoding[next_element_offset]

            if not grammar.char_acceptance_at_element(next_element_offset).get(
                byte, False
            ):
                # if the current byte is not accepted by the current rule, we need to try next rule
                continue

            next_element_offset += num_chars + 1
            new_stack = stk[:-1]
            if grammar.grammar_encoding[next_element_offset]:
                new_stack.append(next_element_offset)
            new_stacks.extend(grammar.advance_stack(tuple(new_stack)))

        if new_stacks:
            check_token_acceptance_in_trie(
                next_trie, new_stacks, grammar, eos_token_id, accepts
            )

    return accepts


if __name__ == "__main__":
    from transformers import AutoTokenizer

    with open("examples/grammars/japanese.ebnf", "r") as file:
        input_text = file.read()
    parsed_grammar = parse_ebnf(input_text)
    parsed_grammar.print()

    tokenizer = AutoTokenizer.from_pretrained("gpt2")

    tokenRecognizer = IncrementalTokenRecognizer(
        grammar_str=input_text, start_rule_name="root", tokenizer=tokenizer
    )

    japanese = "トリーム"  # "こんにちは"
    token_ids = tokenizer.encode(japanese)
    # 13298, 12675, 12045, 254
    stacks = tokenRecognizer._consume_token_ids(
        token_ids, tokenRecognizer.string_recognizer.stacks, as_string=False
    )

    if stacks:
        print("The Japanese input is accepted")
    else:
        print("The Japanese input is not accepted")

    korean = "안녕하세요"
    token_ids = tokenizer.encode(korean)

    try:
        stacks = tokenRecognizer._consume_token_ids(
            token_ids, tokenRecognizer.string_recognizer.stacks, as_string=False
        )
        if stacks:
            print("The Korean input is accepted")
        else:
            print("The Korean input is not accepted")
    except ValueError as e:
        print("The Korean input is not accepted")