Upload chat_template.py

chat_template.py

@@ -0,0 +1,41 @@
+import copy
+
+
+class ChatTemplate:
+    cache = {}
+
+    def __init__(self, model, im_start=r'<|im_start|>', im_end=r'<|im_end|>', nl='\n'):
+        self.model = model
+        self.nl = nl
+        self.im_start = im_start
+        self.im_start_token = model.tokenize(self.im_start.encode('utf-8'), add_bos=False, special=True)
+        self.im_end = im_end
+        self.im_end_nl = model.tokenize((self.im_end + self.nl).encode('utf-8'), add_bos=False, special=True)
+        self.eos = [model._token_eos, self.im_end_nl[0]]
+        self.onenl = [self.im_end_nl[-1]]
+        tmp = model.tokenize(('\r' + self.nl).encode('utf-8'), add_bos=False, special=True)
+        if len(tmp) == 1:
+            self.onenl.append(tmp[0])
+        self.onerl = model.tokenize(b'\r', add_bos=False, special=True)
+        self.nlnl = None
+        tmp = model.tokenize((self.nl + self.nl).encode('utf-8'), add_bos=False, special=True)
+        if len(tmp) == 1:
+            self.nlnl = tmp[0]
+        print('ChatTemplate', self.eos, self.im_end_nl, self.onerl, self.onenl, self.nlnl)
+
+    def _get(self, key: str):
+        if key in self.cache:
+            return copy.deepcopy(self.cache[key])  # 深拷贝一下
+        else:
+            value = self.model.tokenize((self.im_start + key + self.nl).encode('utf-8'), add_bos=False, special=True)
+            self.cache[key] = copy.deepcopy(value)  # 深拷贝一下
+            return value
+
+    def __call__(self, _role, prompt=None):
+        if prompt is None:
+            return self._get(_role)
+        # print(_role, prompt, self.cache)
+        prompt = self.im_start + _role + self.nl + prompt
+        prompt = self.model.tokenize(prompt.encode('utf-8'), add_bos=False, special=True) + self.im_end_nl
+        # print(self.model.str_detokenize(prompt), prompt)
+        return prompt