Update

README.md

@@ -120,11 +120,12 @@ tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-V', trust_remote_code
 model.eval().cuda()
 
 image = Image.open('xx.jpg').convert('RGB')
-question = '请描述一下该图像'
+question = 'What is in the image?'
+msgs = [{'role': 'user', 'content': question}]
 
 res, context, _ = model.chat(
     image=image,
-    question=question,
+    msgs=msgs,
     context=None,
     tokenizer=tokenizer,
     sampling=True,

modeling_minicpmv.py

@@ -235,12 +235,22 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
         return result
 
 
-    def chat(self, image, question, context, tokenizer, vision_hidden_states=None, max_new_tokens=2048, sampling=False, **kwargs):
-        if not context:
-            question = tokenizer.im_start + tokenizer.unk_token * self.config.query_num + tokenizer.im_end + '\n' + question
-            final_input = f'<用户>{question}<AI>'
-        else:
-            final_input = f'{context}<用户>{question}<AI>'
+    def chat(self, image, msgs, context, tokenizer, vision_hidden_states=None, max_new_tokens=2048, sampling=False, **kwargs):
+        if isinstance(msgs, str):
+            msgs = json.loads(msgs)
+        # msgs to prompt
+        prompt = ''
+        for i, msg in enumerate(msgs):
+            role = msg['role']
+            content = msg['content']
+            assert role in ['user', 'assistant']
+            if i == 0:
+                assert role == 'user', 'The role of first msg should be user'
+                content = tokenizer.im_start + tokenizer.unk_token * self.config.query_num + tokenizer.im_end + '\n' + content
+            prompt += '<用户>' if role=='user' else '<AI>'
+            prompt += content
+        prompt += '<AI>'
+        final_input = prompt
 
         if sampling:
             generation_config = {
@@ -268,10 +278,11 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
                 return_vision_hidden_states=True,
                 **generation_config
             )
+        answer = res[0]
+        context = msgs
+        context.append({'role':'assistant', 'content': answer})
 
-        context = final_input + res[0]
-
-        return res[0], context, generation_config
+        return answer, context, generation_config
 
 
 class LlamaTokenizerWrapper(LlamaTokenizer):