#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from transformers import PreTrainedTokenizer
from typing import List
import torch


class ForceLastPhraseLogitsProcessor:
    """
    A logits processor which forces LLMs to use the given phrase before they finalize their answers.
    Most common use cases can be providing references, thanking user with context etc.

    Parameters
    ----------
    phrase (str): The phrase to be generated by LLM before the end of its speech.
    tokenizer (PreTrainedTokenizer): The tokenizer used by the LLM.
    """
    def __init__(self, phrase: str, tokenizer: PreTrainedTokenizer):
        self.eos_token_id = tokenizer.eos_token_id
        self.phrase_tokens = tokenizer.encode(phrase, add_special_tokens=False)
        self._reset()
        self.phrase = phrase
        self.tokenizer = tokenizer

    # LogitsProcessor can contain a clone attribute to deep copy it
    # https://github.com/vllm-project/vllm/blob/19dcc02a72e3ed52e3bf95aae44ea1f40ce42ea0/vllm/sampling_params.py#L537-L550
    def clone(self):
        return ForceLastPhraseLogitsProcessor(self.phrase, self.tokenizer)

    def _reset(self):
        self.index = 0

    def __call__(self, prompt_tokens_ids: List[int], past_token_ids: List[int], scores: torch.Tensor) -> torch.Tensor:
        if not past_token_ids:  # new generation
            self._reset()

        if scores.argmax() == self.eos_token_id and self.index == 0:
            scores[self.phrase_tokens[self.index]] = scores.max() + 1
            self.index += 1
        elif len(self.phrase_tokens) > self.index > 0:
            scores[self.phrase_tokens[self.index]] = scores.max() + 1
            self.index += 1

        return scores

#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from transformers import PreTrainedTokenizer
from typing import List
import torch
from logits_processor_zoo.utils import text_to_token, get_new_line_tokens


class MultipleChoiceLogitsProcessor:
    """
    A logits processor to answer multiple choice questions with one of the choices.
    A multiple choice question is like:
    I am getting a lot of calls during the day. What is more important for me to consider when I buy a new phone?
    0. Camera
    1. Screen resolution
    2. Operating System
    3. Battery
    The goal is to make LLM generate "3" as an answer.


    Parameters
    ----------
    tokenizer (PreTrainedTokenizer): The tokenizer used by the LLM.
    choices (List[str]): List of one character answers like A, B, C, D.
    delimiter (str): One character delimiter that comes after the choices like 1. or 2-.
    boost_first_words (float): Nonzero values add choices' first tokens' logits to boost performance.
                            Especially useful for the models which have difficulty associating the choice with its text.
    """
    def __init__(self, tokenizer: PreTrainedTokenizer, choices: List[str] = None,
                 delimiter: str = ".", boost_first_words: float = 0.0):
        self.tokenizer = tokenizer
        self.choices = choices
        self.delimiter = delimiter
        if choices is None:
            choices = ["1", "2", "3", "4"]

        self.new_line_token = get_new_line_tokens(tokenizer)
        self.delimiter_token = text_to_token(tokenizer, delimiter, last=False)
        self.choice_tokens = [text_to_token(tokenizer, choice, last=False) for choice in choices]
        self.boost_first_words = boost_first_words
        self.very_large_number = 999

    def clone(self):
        return MultipleChoiceLogitsProcessor(self.tokenizer, self.choices, self.delimiter, self.boost_first_words)

    def __call__(self, prompt_tokens_ids: List[int], past_token_ids: List[int], scores: torch.Tensor) -> torch.Tensor:

        if self.boost_first_words:
            choice = 0

            first_tokens = []
            for i in range(len(prompt_tokens_ids) - 3):
                # A choice is like "\nA) hair dryer", where first token is "hair"
                choice_starts = (
                        (prompt_tokens_ids[i] in self.new_line_token) and
                        (prompt_tokens_ids[i + 1] == self.choice_tokens[choice]) and
                        (prompt_tokens_ids[i + 2] == self.delimiter_token)
                )

                if choice_starts:
                    first_tokens.append(prompt_tokens_ids[i + 3])
                    choice += 1

                    if choice >= len(self.choice_tokens):
                        break

            scores[self.choice_tokens[:len(first_tokens)]] += self.boost_first_words * scores[first_tokens]

        scores[self.choice_tokens] += self.very_large_number
        return scores

#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from transformers import PreTrainedTokenizer
from typing import List
import torch
from logits_processor_zoo.utils import text_to_token


class TriggerPhraseLogitsProcessor:
    """
    A logits processor which triggers phrases when it encounters a given token.

    Parameters
    ----------
    phrase (str): The phrase to be generated by LLM when it encounters the trigger token.
    trigger_token_phrase (str): One token phrase in string to trigger phrases.
    tokenizer (PreTrainedTokenizer): The tokenizer used by the LLM.
    trigger_count (int): How many times the phrase will be triggered.
    trigger_after (bool): Whether the phrase is written after the trigger token or instead of the trigger token.
    """
    def __init__(self, phrase: str, trigger_token_phrase: str, tokenizer: PreTrainedTokenizer, trigger_count: int = 1,
                 trigger_after: bool = False):
        self.phrase = phrase
        self.trigger_token_phrase = trigger_token_phrase
        self.tokenizer = tokenizer
        self.trigger_count = trigger_count
        self.trigger_token = text_to_token(tokenizer, trigger_token_phrase, last=False)
        self.phrase_tokens = tokenizer.encode(phrase, add_special_tokens=False)
        self.initial_trigger_count = trigger_count
        self.trigger_after = trigger_after
        self.very_large_number = 999
        self._reset()

    def clone(self):
        return TriggerPhraseLogitsProcessor(self.phrase, self.trigger_token_phrase, self.tokenizer,
                                            self.initial_trigger_count, self.trigger_after)

    def _reset(self):
        self.index = -1
        self.trigger_count = self.initial_trigger_count

    def __call__(self, prompt_tokens_ids: List[int], past_token_ids: List[int], scores: torch.Tensor) -> torch.Tensor:
        if not past_token_ids:  # new generation
            self._reset()

        if self.trigger_count <= 0:
            return scores

        if scores.argmax() == self.trigger_token and self.index == -1:
            self.index = 0
            if not self.trigger_after:
                scores[self.phrase_tokens[self.index]] = scores.max() + self.very_large_number
                self.index += 1
        elif len(self.phrase_tokens) > self.index >= 0:
            scores[self.phrase_tokens[self.index]] = scores.max() + self.very_large_number
            self.index += 1

        if len(self.phrase_tokens) == self.index:  # phrase completed, reset for next trigger
            self.index = -1
            self.trigger_count -= 1

        return scores

Create a logits processor that calls a custom function with parameters (completion, last_token). You can come up with your own format for the function return note that this is a library so we want to be the interface through which the user can return a forced/suggested token be very easy to use and understand try to make it as comprehensive as possible. examples are provided above