File size: 1,132 Bytes
c6be546 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | """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
|