How to train with multi images?

[goodstudent9]

Hello,
I haved used VILA-LITE model to do my own multi images tasks. But that seems like the grpo training code doesn't support multi images input? I don't know if I am correct. From the code, we can see that it only get the first image value? And there is some error that message_key is used before aaignment.
How can I solved this problem?
`def _get_messages_vila(example: Dict[str, Any],
prompt_key: str = "prompt",
image_key: str = "images",
image_dir: Optional[str] = None,
video_key: str = "videos",
video_dir: str = None,) -> Dict[str, Any]:
if video_key in example:
vision_key = "video"
vision_value = example[video_key]
if video_dir is not None and isinstance(vision_value, str): # image paths
vision_value = os.path.join(video_dir, vision_value)
message_key = "video"
question_template = QUESTION_TEMPLATE_VIDEO
elif image_key in example:
vision_key = "image"
vision_value = example[image_key][0]
if isinstance(vision_value, ImageObject):
message_key = "image_pil"
elif isinstance(vision_value, str):
vision_key = "image"
else:
raise ValueError("Unknown image type", vision_value)
question_template = QUESTION_TEMPLATE_IMAGE
else:
raise ValueError("Unsupported VILA for text only.")

messages = [{"role": "user", "content": "<%s>" % vision_key + example[prompt_key]}]
prompt = question_template.format(question=messages[-1]['content'].replace("<%s>" % vision_key, ""))
messages[-1]['content'] = [
    {"type": vision_key, message_key: vision_value},
    {"type": "text", "text": prompt},
]
return messages, prompt`