"""Inference utilities vendored from xtuner (stop criteria + prompt template)."""

from transformers import StoppingCriteria, StoppingCriteriaList


PHI3_CHAT_TEMPLATE = {
    'SYSTEM': '<|system|>\n{system}<|end|>\n',
    'INSTRUCTION': '<|user|>\n{input}<|end|>\n<|assistant|>\n',
    'SUFFIX': '<|end|>',
    'SUFFIX_AS_EOS': True,
    'SEP': '\n',
    'STOP_WORDS': ['<|end|>'],
}


class StopWordStoppingCriteria(StoppingCriteria):
    """Stop generation once the decoded tail matches a stop word."""

    def __init__(self, tokenizer, stop_word):
        self.tokenizer = tokenizer
        self.stop_word = stop_word
        self.length = len(stop_word)

    def __call__(self, input_ids, *args, **kwargs) -> bool:
        current_text = self.tokenizer.decode(input_ids[0])
        current_text = current_text.replace('\r', '').replace('\n', '')
        return current_text[-self.length:] == self.stop_word


def get_stop_criteria(tokenizer, stop_words=()):
    stop_criteria = StoppingCriteriaList()
    for word in stop_words:
        stop_criteria.append(StopWordStoppingCriteria(tokenizer, word))
    return stop_criteria