Refactor chat template and support accurate name matching.

examples/vl/qwen_model.py

@@ -19,20 +19,12 @@ def __init__(self,
                  top_p=0.3,
                  top_k=None,
                  temperature=1.0,
-                 im_start='<|im_start|>',
-                 im_end='<|im_end|>',
-                 system='You are a helpful assistant.',
-                 stop_words=['<|im_end|>'],
                  **kwargs):
         super().__init__(**kwargs)
         self.session_len = session_len
         self.top_p = top_p
         self.top_k = top_k
         self.temperature = temperature
-        self.im_start = im_start
-        self.im_end = im_end
-        self.system = system
-        self.stop_words = stop_words
 
     def _concat_image_info(self, prompt):
         """Append image placeholder."""
@@ -45,27 +37,30 @@ def _concat_image_info(self, prompt):
         prompt = res + prompt
         return prompt
 
-    def decorate_prompt(self, prompt, sequence_start=True):
+    def get_prompt(self, prompt, sequence_start=True):
         """Apply chat template to prompt."""
         prompt = self._concat_image_info(prompt)
-        return super().decorate_prompt(prompt, sequence_start)
+        return super().get_prompt(prompt, sequence_start)
 
     def messages2prompt(self, messages, sequence_start=True):
         """Apply chat template to history."""
         if isinstance(messages, str) or isinstance(messages[0], str):
-            return self.decorate_prompt(messages, sequence_start)
-        system, users, assistants = self._translate_messages(messages)
-        ret = f'{self.im_start}system\n{system}{self.im_end}'
-        for user, assistant in zip(users, assistants):
-            if not isinstance(user):
-                user = [user[0]['text'], len(user) - 1]
-                user = self._concat_image_info(user)
-            if assistant:
-                ret += f'\n{self.im_start}user\n{user}{self.im_end}' \
-                       f'\n{self.im_start}assistant\n{assistant}'
-            else:
-                ret += f'\n{self.im_start}user\n{user}{self.im_end}' \
-                       f'\n{self.im_start}assistant\n'
+            return self.get_prompt(messages, sequence_start)
+        box_map = dict(user=self.user,
+                       assistant=self.assistant,
+                       system=self.system)
+        eox_map = dict(user=self.eoh,
+                       assistant=self.eoa + self.stop_word_suffix,
+                       system=self.eosys)
+        ret = ''
+        for message in messages:
+            role = message['role']
+            content = message['content']
+            if role == 'user' and not isinstance(content):
+                content = [content[0]['text'], len(content) - 1]
+                content = self._concat_image_info(content)
+            ret += f'{box_map[role]}{content}{eox_map[role]}'
+        ret += f'{self.assistant}'
         return ret
 
 
@@ -134,8 +129,8 @@ def prepare_query(self, query, sequence_start=True):
         image_paths = []
         if not isinstance(query, str):
             query, image_paths = query[0], query[1:]
-        decorate_text = self.decorator.decorate_prompt(
-            (query, len(image_paths)), sequence_start)
+        decorate_text = self.decorator.get_prompt((query, len(image_paths)),
+                                                  sequence_start)
         return self._to_inputs(decorate_text, image_paths, sequence_start)
 
     def prepare_message(self, messages):

examples/vl/xcomposer_model.py

@@ -9,7 +9,7 @@
 from huggingface_hub import snapshot_download
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
 
-from lmdeploy.model import MODELS, BaseModel
+from lmdeploy.model import MODELS, BaseChatTemplate
 
 meta_instruction = """meta instruction
 You are an AI assistant whose name is 浦语.
@@ -20,20 +20,20 @@
 
 
 @MODELS.register_module(name='internlm-xcomposer-7b')
-class InternLMXComposerTemplate(BaseModel):
+class InternLMXComposerTemplate(BaseChatTemplate):
     """Internlm xcomposer chat template."""
 
     def __init__(self,
-                 system=meta_instruction,
-                 user='<|User|>:',
-                 assistant='<|Bot|>:',
+                 meta_instruction=meta_instruction,
+                 user=' <|User|>: ',
+                 assistant=' <|Bot|>: ',
                  eoh='<TOKENS_UNUSED_0>',
                  eoa='<TOKENS_UNUSED_1>',
                  stop_words=['<TOKENS_UNUSED_0>', '<TOKENS_UNUSED_1>'],
                  image_placeholder='<Img><ImageHere></Img>',
                  **kwargs):
         super().__init__(**kwargs)
-        self.system = system
+        self.meta_instruction = meta_instruction
         self.user = user
         self.assistant = assistant
         self.eoh = eoh
@@ -51,31 +51,32 @@ def _concat_image_info(self, prompt):
             prompt = f'{self.image_placeholder}{prompt}'
         return prompt
 
-    def decorate_prompt(self, prompt, sequence_start=True):
+    def get_prompt(self, prompt, sequence_start=True):
         """Apply chat template to prompt."""
         prompt = self._concat_image_info(prompt)
-        if sequence_start:
-            return f'{self.system} {self.user} {prompt}{self.eoh} {self.assistant}'  # noqa
-        else:
-            return f' {self.user} {prompt}{self.eoh} {self.assistant}'
+        return super().get_prompt(prompt, sequence_start)
 
     def messages2prompt(self, messages, sequence_start=True):
         """Apply chat template to history."""
         if isinstance(messages, str) or isinstance(messages[0], str):
-            return self.decorate_prompt(messages, sequence_start)
-        system, users, assistants = self._translate_messages(messages)
-        system = self.system if not system else system
-        ret = system
-        for user, assistant in zip(users, assistants):
-            if not isinstance(user, str):
-                assert isinstance(user, Sequence)
-                assert all(isinstance(item, dict) for item in user)
-                user = [user[0]['text'], len(user) - 1]
-            user = self._concat_image_info(user)
-            if assistant:
-                ret += f' {self.user} {user}{self.eoh} {self.assistant} {assistant}{self.eoa}'  # noqa
-            else:
-                ret += f' {self.user} {user}{self.eoh} {self.assistant}'
+            return self.get_prompt(messages, sequence_start)
+        box_map = dict(user=self.user,
+                       assistant=self.assistant,
+                       system=self.system)
+        eox_map = dict(user=self.eoh,
+                       assistant=self.eoa + self.stop_word_suffix,
+                       system=self.eosys)
+        ret = ''
+        for message in messages:
+            role = message['role']
+            content = message['content']
+            if role == 'user' and not isinstance(content, str):
+                assert isinstance(content, Sequence)
+                assert all(isinstance(item, dict) for item in content)
+                content = [content[0]['text'], len(content) - 1]
+            content = self._concat_image_info(content)
+            ret += f'{box_map[role]}{content}{eox_map[role]}'
+        ret += f'{self.assistant}'
         return ret
 
 
@@ -155,8 +156,7 @@ def prepare_query(self, query, sequence_start=True):
             if len(image_paths) > 1:
                 print('does not support multiple images, use last one.')
                 image_paths = image_paths[-1:]
-        decorate_text = self.decorator.decorate_prompt(
-            (query, len(image_paths)))
+        decorate_text = self.decorator.get_prompt((query, len(image_paths)))
         return self._to_inputs(decorate_text, image_paths, sequence_start)
 
     def prepare_message(self, messages):