add UT

Author: lvhan028

lmdeploy/vl/model/internvl3_hf.py

@@ -147,6 +147,30 @@ def forward(self, messages: List[Dict], max_batch_size: int = 1) -> List[Dict]:
         messages.append(dict(role='forward', content=outputs))
         return messages
 
+    def proc_internvl_hf_messages(self, content: List[Dict]):
+        """Process the content list of role 'user' for InternVL HF models."""
+        res = []
+        for item in content:
+            if item['type'] == 'text':
+                res.append(item['text'])
+            elif item['type'] in ['image', 'image_url']:
+                res.append(f'{self.image_token}\n')
+            else:
+                raise ValueError(f'Unsupported message type: {item["type"]}')
+        return ''.join(res)
+
+    def proc_interns1_messages(self, content: List[Dict]):
+        """Process the content list of role 'user' for InternS1 models."""
+        res = []
+        for item in content:
+            if item['type'] == 'text':
+                res.append(item['text'])
+            elif item['type'] in ['image', 'image_url']:
+                res.append(f'{self.image_token}')
+            else:
+                raise ValueError(f'Unsupported message type: {item["type"]}')
+        return '\n'.join(res)
+
     def proc_messages(
         self,
         messages,
@@ -162,15 +186,9 @@ def proc_messages(
                 continue
             role, content = message['role'], message['content']
             if role == 'user' and isinstance(content, List):
-                _content = []
-                for item in content:
-                    if item['type'] == 'text':
-                        _content.append(item['text'])
-                    elif item['type'] in ['image', 'image_url']:
-                        _content.append(self.image_token)
-                    else:
-                        raise ValueError(f'Unsupported message type: {item["type"]}')
-                message = dict(role=role, content='\n'.join(_content))
+                content = (self.proc_internvl_hf_messages(content)
+                           if self.arch == 'InternVLForConditionalGeneration' else self.proc_interns1_messages(content))
+                message = dict(role=role, content=content)
                 prompt_messages.append(message)
             else:
                 prompt_messages.append(message)

tests/test_lmdeploy/test_vl/test_internvl3_hf.py

@@ -0,0 +1,33 @@
+import pytest
+
+from lmdeploy.model import HFChatTemplate
+from lmdeploy.vl.model.internvl3_hf import InternVL3VisionModel
+
+TEST_MODELS = ['OpenGVLab/InternVL3_5-8B-HF', 'internlm/Intern-S1-mini']
+
+
+@pytest.fixture(scope='module')
+def mock_messages():
+    return [
+        dict(role='user',
+             content=[
+                 dict(type='text', text='Describe the following images in detail'),
+                 dict(type='image', url=dict(url='http://images.cocodataset.org/val2017/000000039769.jpg')),
+                 dict(type='image', url=dict(url='http://images.cocodataset.org/val2017/000000039769.jpg')),
+                 dict(type='text', text='How many cats are there in total?')
+             ]),
+    ]
+
+
+def test_proc_messages(mock_messages):
+    for model_path in TEST_MODELS:
+        vision_model = InternVL3VisionModel(model_path=model_path, with_llm=False)
+        vision_model.build_preprocessor()
+        reference = vision_model.processor.apply_chat_template(mock_messages,
+                                                               add_generation_prompt=True,
+                                                               tokenize=False,
+                                                               return_dict=True)
+        chat_template = HFChatTemplate(model_path=model_path)
+        vision_model.proc_messages(mock_messages, chat_template, sequence_start=True)
+        prompt, _ = vision_model.proc_messages(mock_messages, chat_template, sequence_start=True)
+        assert prompt == reference