logits_processor.py

from __future__ import annotations

import copy
import logging
import math
import time
from types import SimpleNamespace

import torch
from tqdm.auto import tqdm
from transformers import CodeLlamaTokenizer, CodeLlamaTokenizerFast, GenerationConfig, LlamaTokenizer, \
    LlamaTokenizerFast, LogitsProcessor, PreTrainedTokenizer, RobertaTokenizer, RobertaTokenizerFast

from generation_rules import GenerationRule, GenerationRuleset
from openapi_utils import spec_to_ruleset

logging.basicConfig()
logger = logging.getLogger(__name__)
logger.setLevel(logging.WARNING)

# Cancel the constrained generation after 15 min
TIMEOUT = 15 * 60


class ConstrainedDecoder(LogitsProcessor):
    """
    [`LogitsProcessor`] that enables constrained decoding based on a generation ruleset.
    :param generation_config: The GenerationConfig used
    :param tokenizer: The tokenizer used
    :param ruleset: A list of rules for generation
    :param completion_prefix: Suffix of the prompt that should be treated as if it was generated by the model
        (important if the prompt includes some starter code that could trigger a rule)
    :param explain: Explain every generation step (has some performance overhead)
    """

    def __init__(self, generation_config: GenerationConfig, tokenizer: PreTrainedTokenizer, ruleset: GenerationRuleset,
                 completion_prefix: str = "", explain: bool = False) -> None:
        super().__init__()

        self._tokenizer = tokenizer
        self._ruleset_template = ruleset
        self._completion_prefix = completion_prefix
        self._explain = explain

        self._num_beams = generation_config.num_beams
        self._top_k = generation_config.top_k if generation_config.do_sample else 1

        if self._num_beams > 1:
            raise ValueError(f"Beam search is not yet supported; got {generation_config.num_beams = }")

        if isinstance(self._tokenizer, LlamaTokenizerFast) or isinstance(self._tokenizer, LlamaTokenizer) or \
                isinstance(self._tokenizer, CodeLlamaTokenizerFast) or isinstance(self._tokenizer, CodeLlamaTokenizer):
            self._is_llama_tokenizer = True
            self._token_0_length = len(self._tokenizer.decode(0))
        else:
            self._is_llama_tokenizer = False

        if self._explain and logger.getEffectiveLevel() > logging.INFO:
            logger.warning(f"Setting logging level to INFO to be able to explain the constrained decoding process")
            logger.setLevel(logging.INFO)

        self._all_ids = list(range(len(self._tokenizer.get_vocab())))
        self._normal_ids_mask = [id not in self._tokenizer.all_special_ids for id in self._all_ids]

        self._states = []  # to handle multiple parallel generations, we store multiple instances of internal state

        self._ruleset_template.reset()  # just in case the ruleset was cached and has been used before
        self._ruleset_template.update(self._completion_prefix)  # fast-forward the internal state of the ruleset

    def reset(self, completion_prefix: str | None = None) -> None:
        """
        Reset the internal state of this ConstrainedDecoder and its GenerationRuleset.
        :param completion_prefix: Optional *new* completion prefix (otherwise the previous prefix is retained)
        """
        if completion_prefix is not None and completion_prefix != self._completion_prefix:
            self._completion_prefix = completion_prefix
            self._ruleset_template.reset()
            self._ruleset_template.update(self._completion_prefix)

        self._states = []

    def _get_state(self, state_id: int) -> SimpleNamespace:
        """
        Get the state with this ID, creating a new one if none exists.
        :param state_id: The ID of the state to retrieve
        :return: A container of all stateful data
        """
        if state_id < len(self._states):
            return self._states[state_id]

        logger.debug(f"Creating new state with ID {state_id}")
        assert state_id == len(self._states)
        state = SimpleNamespace(
            ruleset=copy.deepcopy(self._ruleset_template),
            generation_step=0,
            start_time=0,
            completion=self._completion_prefix,
            timed_out=False,
            constraints_unsatisfiable=False,
            interventions=0,
            last_most_likely_id=-1,
        )
        self._states.append(state)
        return state

    # noinspection PyTypeChecker
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        """
        Implementation and docstring copied from :class:`PrefixConstrainedLogitsProcessor`.
        :param input_ids: (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
            Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids)
        :param scores: (`torch.FloatTensor` of shape `(batch_size, config.vocab_size)`):
            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam
            search or log softmax for each vocabulary token when using beam search
        :return: `torch.FloatTensor` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.
        """
        mask = torch.full_like(scores, -math.inf)
        batch_size = input_ids.shape[0] // self._num_beams

        for batch_id in range(batch_size):
            for beam_id in range(self._num_beams):
                sent = input_ids[batch_id * self._num_beams + beam_id]
                prefix_allowed_tokens = self._prefix_allowed_tokens_fn(batch_id, sent, scores[batch_id])
                if len(prefix_allowed_tokens) == 0:
                    raise ValueError(
                        f"`prefix_allowed_tokens_fn` returned an empty list for batch ID {batch_id}."
                        f"This means that the constraint is unsatisfiable. Please check your implementation"
                        f"of `prefix_allowed_tokens_fn` "
                    )
                mask[batch_id * self._num_beams + beam_id, prefix_allowed_tokens] = 0

        scores_processed = scores + mask
        return scores_processed

    def _prefix_allowed_tokens_fn(self, batch_id: int, input_ids: torch.LongTensor, scores: torch.FloatTensor) \
            -> list[int]:
        """
        This function constraints the beam search to allowed tokens only at each step. This function takes 2
        arguments `inputs_ids` and the batch ID `batch_id`. It has to return a list with the allowed tokens for the
        next generation step conditioned on the previously generated tokens `inputs_ids` and the batch ID
        `batch_id`.
        :param batch_id: The ID of the current batch
        :param input_ids: The IDs of the tokens generated so far
        :param scores: The logits for the next token prediction
        :return: A list of allowed token IDs
        """
        state = self._get_state(batch_id)

        state.generation_step += 1
        logger.debug(f"Batch #{batch_id}, step #{state.generation_step}")

        if state.generation_step == 1:
            state.start_time = time.process_time()

        # Depending on the tokenizer, the first few tokens are always the same
        if state.generation_step <= 3:
            if isinstance(self._tokenizer, RobertaTokenizerFast) or isinstance(self._tokenizer, RobertaTokenizer):
                # The first 3 tokens are always <pad>, <s>, <extra_id_0>
                if state.generation_step == 1:
                    return [self._tokenizer.pad_token_id]
                elif state.generation_step == 2:
                    return [self._tokenizer.bos_token_id]
                elif state.generation_step == 3:
                    state.last_most_likely_id = scores.argmax().item()
                    return self._tokenizer.additional_special_tokens_ids
            elif state.generation_step == 1:
                logger.info(f"Assuming {type(self._tokenizer).__name__} doesn't need a special starting sequence")

        # Find out which token was selected in the previous iteration and update the internal state accordingly
        if state.generation_step >= 2:
            last_selected_id = input_ids[-1].item()
            last_token = self._decode_token(last_selected_id)

            if self._is_normal_id(last_selected_id):
                state.completion += last_token

            if self._explain:
                logger.info(f"Last generated token: {repr(last_token)}")
                if self._top_k == 1 and last_selected_id != state.last_most_likely_id:
                    assert state.ruleset.has_active_rules(), "Intervention shouldn't be possible if no rule is active"
                    state.interventions += 1
                    last_most_likely_token = self._decode_token(state.last_most_likely_id)
                    logger.info(f"Most likely token would have been {repr(last_most_likely_token)}")
                    logger.info(f"This was intervention #{state.interventions}")

        if self._explain:
            logger.info(f"Completion so far:\n{state.completion}")
            match = state.ruleset.match_whole_code(state.completion, excluded=["funnel"], partial=True)
            assert match, "Completion doesn't satisfy the constraints - this should be impossible"
            logger.info(f"Completion satisfies the constraints:\n{match}")

            state.last_most_likely_id = scores.argmax().item()

        # Check if any generation rules become active or inactive
        has_active_rules = state.ruleset.update(state.completion)

        logger.debug("Ruleset has been updated")
        if self._explain:
            state.ruleset.print_state(completion=state.completion, active_only=True, logger=logger)

        # If the generation is currently unconstrained, we can skip the next steps
        if not has_active_rules:
            return self._all_ids

        logger.debug(f"Searching for {self._top_k} valid token(s) ...")
        allowed_ids = []
        ids = torch.argsort(scores, descending=True)
        try:
            for id in tqdm(ids, disable=logger.getEffectiveLevel() > logging.DEBUG):
                if self._is_normal_id(id):
                    id = id.item()
                    token = self._decode_token(id)
                    remaining_time = max(0.0, TIMEOUT - (time.process_time() - state.start_time))
                    if state.ruleset.is_valid_continuation(state.completion + token, timeout=remaining_time):
                        allowed_ids.append(id)
                        logger.debug(f"Found valid token {repr(token)}")
                        if len(allowed_ids) >= self._top_k:
                            return allowed_ids

            if allowed_ids:
                logger.debug(f"Found only {len(allowed_ids)} valid tokens")
                return allowed_ids

            logger.warning("Constraints unsatisfiable - terminating sequence")
            state.constraints_unsatisfiable = True
            return [self._tokenizer.eos_token_id]

        except TimeoutError:
            logger.warning("Timeout - terminating sequence")
            state.timed_out = True
            return [self._tokenizer.eos_token_id]

    def _decode_token(self, id: int) -> str:
        """
        Convert a token ID to a token string, handling some edge cases.
        :param id: The token ID to decode
        :return: The corresponding token string
        """
        if self._is_llama_tokenizer:
            # This detokenizer removes leading whitespace. As a workaround, we temporarily prepend some token.
            return self._tokenizer.decode([0, id])[self._token_0_length:]
        else:
            return self._tokenizer.decode(id)

    def _is_normal_id(self, id: int | torch.Tensor) -> bool:
        """
        Check if the given token ID is normal, i.e., not a special one.
        :param id: The ID to check
        :return: Whether the ID is normal
        """
        return False if id >= len(self._normal_ids_mask) else self._normal_ids_mask[id]

    def constraints_unsatisfiable(self, state_id: int) -> bool:
        """
        Get the ``constraints_unsatisfiable`` flag for the given ID.
        :param state_id: ID of the state/batch
        :return: ``constraints_unsatisfiable``
        """
        return self._states[state_id].constraints_unsatisfiable

    def timed_out(self, state_id: int) -> bool:
        """
        Get the ``timed_out`` flag for the given ID.
        :param state_id: ID of the state/batch
        :return: ``timed_out``
        """
        return self._states[state_id].timed_out


class SimpleRegexDecoder(ConstrainedDecoder):
    """
    Constrained decoding from a single regex.
    :param generation_config: The GenerationConfig used
    :param tokenizer: The tokenizer used
    :param regex: The regex
    """

    def __init__(self, generation_config: GenerationConfig, tokenizer: PreTrainedTokenizer, regex: str,
                 **kwargs) -> None:
        super().__init__(generation_config, tokenizer, GenerationRuleset([GenerationRule(None, None, regex)]), **kwargs)

        self._regex = regex


class OpenApiDecoder(ConstrainedDecoder):
    """
    Constrained decoding from an OpenAPI specification.
    :param generation_config: The GenerationConfig used
    :param tokenizer: The tokenizer used
    :param spec_path: Path to the OpenAPI specification
    """

    def __init__(self, generation_config: GenerationConfig, tokenizer: PreTrainedTokenizer, spec_path: str,
                 **kwargs) -> None:
        super().__init__(generation_config, tokenizer, spec_to_ruleset(spec_path), **kwargs)

        self._spec_path = spec_path