diff --git a/VL/llava/conversation.py b/VL/llava/conversation.py index ad6bf539..587f10f7 100644 --- a/VL/llava/conversation.py +++ b/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. " diff --git a/VL/llava/test_user_content.py b/VL/llava/test_user_content.py new file mode 100644 index 00000000..a6511f16 --- /dev/null +++ b/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"}], + } + ], + "", + ) + 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"}, + ], + "", + ) + self.assertEqual(query, "\nhello\nworld") + self.assertEqual(url, "a.png") + + query, url = conv_mod.render_user_content("hi there", "") + self.assertEqual(query, "hi there") + self.assertEqual(url, "") + + +if __name__ == "__main__": + unittest.main() diff --git a/VL/openai_api.py b/VL/openai_api.py index 6a9b2d35..29bc1b9c 100644 --- a/VL/openai_api.py +++ b/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()