Keep a text part when the message has no image

VL/llava/conversation.py

@@ -181,6 +181,42 @@ def dict(self):
         }
 
 
+def render_user_content(content, image_token):
+    if isinstance(content, str):
+        return content, ""
+    num_images = 0
+    image_url = ""
+    texts = []
+    for item in content:
+        if item["type"] == "image_url":
+            num_images += 1
+            image_url = item["image_url"]["url"]
+        elif item["type"] == "text":
+            texts.append(item["text"])
+    prompt = "\n".join(texts)
+    if num_images:
+        return image_token * num_images + "\n" + prompt, image_url
+    return prompt, ""
+
+
+def append_chat_messages(conv, messages, image_token):
+    image_file = ""
+    has_image = False
+    for message in messages:
+        role = message["role"]
+        content = message["content"]
+        if role != "user":
+            conv.append_message(conv.roles[1], content)
+            continue
+        query, message_image = render_user_content(content, image_token)
+        if message_image:
+            image_file = message_image
+            has_image = True
+        if query or message_image:
+            conv.append_message(conv.roles[0], query)
+    return image_file, has_image
+
+
 mm_default_conv = Conversation(
     system="This is a chat between an inquisitive human and an AI assistant. "
     "Assume the role of the AI assistant. "

VL/llava/test_user_content.py

@@ -0,0 +1,51 @@
+import importlib.util
+import unittest
+from pathlib import Path
+
+
+def load_conv():
+    path = Path(__file__).resolve().parent / "conversation.py"
+    spec = importlib.util.spec_from_file_location("yi_conversation", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+class UserContentTest(unittest.TestCase):
+    def test_a_text_part_is_kept_without_an_image(self):
+        conv_mod = load_conv()
+        conv = conv_mod.conv_templates["mm_default"].copy()
+        image_file, has_image = conv_mod.append_chat_messages(
+            conv,
+            [
+                {
+                    "role": "user",
+                    "content": [{"type": "text", "text": "where is the cat"}],
+                }
+            ],
+            "<image>",
+        )
+        conv.append_message(conv.roles[1], None)
+        prompt = conv.get_prompt()
+        self.assertIn("where is the cat", prompt)
+        self.assertEqual(image_file, "")
+        self.assertFalse(has_image)
+
+        query, url = conv_mod.render_user_content(
+            [
+                {"type": "text", "text": "hello"},
+                {"type": "image_url", "image_url": {"url": "a.png"}},
+                {"type": "text", "text": "world"},
+            ],
+            "<image>",
+        )
+        self.assertEqual(query, "<image>\nhello\nworld")
+        self.assertEqual(url, "a.png")
+
+        query, url = conv_mod.render_user_content("hi there", "<image>")
+        self.assertEqual(query, "hi there")
+        self.assertEqual(url, "")
+
+
+if __name__ == "__main__":
+    unittest.main()

VL/openai_api.py

@@ -34,7 +34,7 @@
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
-from llava.conversation import conv_templates
+from llava.conversation import append_chat_messages, conv_templates
 from llava.mm_utils import (
     expand2square,
     get_model_name_from_path,
@@ -498,29 +498,9 @@ def convert_to_inputs(
         """
         conv = conv_templates["mm_default"].copy()
         stop_str = conv.sep
-        image_file = ""
-        has_image = False
-        for message in prompt_or_messages:
-            role = message["role"]
-            content = message["content"]
-            if role != "user":
-                conv.append_message(conv.roles[1], content)
-                continue
-            elif isinstance(content, str):
-                conv.append_message(conv.roles[0], content)
-                continue
-
-            num_images = 0
-            for item in content:
-                if item["type"] == "image_url":
-                    num_images += 1
-                    image_file = item["image_url"]["url"]
-                    has_image = True
-                elif item["type"] == "text":
-                    prompt = item["text"]
-            if image_file != "" and image_file != None:
-                query = DEFAULT_IMAGE_TOKEN * num_images + "\n" + prompt
-                conv.append_message(conv.roles[0], query)
+        image_file, has_image = append_chat_messages(
+            conv, prompt_or_messages, DEFAULT_IMAGE_TOKEN
+        )
 
         conv.append_message(conv.roles[1], None)
         prompt = conv.get_prompt()