| """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 |
|
|
