File size: 1,614 Bytes
01cf98b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 | import re
def build_messages(sys, user, images, raw_res, num_docs=None, index=None, titles=None, instruction=None):
messages = []
user_list = []
if sys is not None:
messages.append({"role": "system", "content": sys})
for segment in re.split("(<image>|<video>|<index>|<title>|<n_images>|<instruction>)", user):
if segment == "<image>":
image_dict = images.pop(0)
user_list.append({"type": "image", 'image': image_dict.get('image'), 'bbox': image_dict.get('bbox', None), 'dynamic_crop': image_dict.get('dynamic_crop', False), 'max_pixels': image_dict.get('max_pixels', None)})
elif segment == "<index>":
user_list.append({"type": "text", "text": str(index)})
elif segment == "<title>":
title = titles.pop(0)
user_list.append({"type": "text", "text": title})
elif segment == "<n_images>":
for i in range(num_docs):
image_dict = images.pop(0)
user_list.append({"type": "image", 'image': image_dict.get('image'), 'bbox': image_dict.get('bbox', None), 'dynamic_crop': image_dict.get('dynamic_crop', False), 'max_pixels': image_dict.get('max_pixels', None)})
elif segment == "<instruction>":
user_list.append({"type": "text", "text": instruction or ""})
else:
user_list.append({"type": "text", "text": segment})
messages.append({"role": "user", "content": user_list})
if raw_res is not None:
messages.append({"role": "assistant", "content": [{'type': 'text', 'text': raw_res}]})
return messages
|