Limour/llama-python-streamingllm

1import copy2 3 4class ChatTemplate:5    cache = {}6    roles = set()7 8    def __init__(self, model, im_start=r'<|im_start|>', im_end=r'<|im_end|>', nl='\n'):9        self.model = model10        self.nl = nl11        self.im_start = im_start12        self.im_start_token = model.tokenize(self.im_start.encode('utf-8'), add_bos=False, special=True)13        self.im_end = im_end14        self.im_end_nl = model.tokenize((self.im_end + self.nl).encode('utf-8'), add_bos=False, special=True)15        self.eos = [model._token_eos, self.im_end_nl[0]]16        self.onenl = [self.im_end_nl[-1]]17        tmp = model.tokenize(('\r' + self.nl).encode('utf-8'), add_bos=False, special=True)18        if len(tmp) == 1:19            self.onenl.append(tmp[0])20        self.onerl = model.tokenize(b'\r', add_bos=False, special=True)21        self.nlnl = None22        tmp = model.tokenize((self.nl + self.nl).encode('utf-8'), add_bos=False, special=True)23        if len(tmp) == 1:24            self.nlnl = tmp[0]25        print('ChatTemplate', self.eos, self.im_end_nl, self.onerl, self.onenl, self.nlnl)26 27    def _get(self, key: str):28        if key in self.cache:29            return copy.deepcopy(self.cache[key])  # 深拷贝一下30        else:31            value = self.model.tokenize((self.im_start + key + self.nl).encode('utf-8'), add_bos=False, special=True)32            self.cache[key] = copy.deepcopy(value)  # 深拷贝一下33            return value34 35    def _add_role(self, _role):36        if _role:37            self.roles.add('\n' + _role)38 39    def eos_in_role(self, history: str, t_bot):40        if not (history.endswith('\n') or history.endswith('\r')):41            return 042        tmp = history.rstrip()43        for _role in self.roles:44            if tmp.endswith(_role):45                n = len(t_bot)46                for i in range(1, n):  # 找出需要弃置的tokens长度47                    tmp = self.model.str_detokenize(t_bot[n - i:])48                    if tmp.rstrip().endswith(_role):49                        print('eos_in_role', t_bot[n - i:], repr(tmp))50                        return i51                print('eos_in_role missing')52                break53        return 054 55    def eos_in_nlnl(self, history: str, t_bot):56        if not (history.endswith('\n\n') or history.endswith('\n\r\n')):57            return 058        n = len(t_bot)59        for i in range(1, n):  # 找出需要弃置的tokens长度60            tmp = self.model.str_detokenize(t_bot[n - i:])61            if tmp.endswith('\n\n') or tmp.endswith('\n\r\n'):62                if tmp.startswith(']'):  # 避免误判63                    return 064                print('eos_in_nlnl', t_bot[n - i:], repr(tmp))65                return i66        print('eos_in_nlnl missing')67        return 068 69    def __call__(self, _role, prompt=None):70        self._add_role(_role)71        if prompt is None:72            return self._get(_role)73        # print(_role, prompt, self.cache)74        prompt = self.im_start + _role + self.nl + prompt75        prompt = self.model.tokenize(prompt.encode('utf-8'), add_bos=False, special=True) + self.im_end_nl76        # print(self.model.str_detokenize(prompt), prompt)77        return prompt78