import unsloth
import trl
import transformers
import torch

print(f"Unsloth version: {unsloth.__version__}")
print(f"TRL version: {trl.__version__}")
print(f"Transformers version: {transformers.__version__}")
print(f"PyTorch version: {torch.__version__}")

def make_conversation(example):
    # Define placeholder constants if they are not defined globally

    # The user's text prompt
    text_content = (example['overall_prompt'])

    image_1 = Image.open(example['img_1_path']).convert("RGB")
    image_2 = Image.open(example['img_2_path']).convert("RGB")

    image_list = [image_1, image_2]

    # Construct the prompt in the desired multi-modal format
    prompt = [
        {
            "role": "user",
            "content": [
                {"type": "image"},  # Placeholder for the image 1
                {"type": "image"},  # Placeholder for the image 2
                {"type": "text", "text": text_content},  # The text part of the prompt
            ],
        },
    ]

    # The actual image data is kept separate for the processor
    return {"prompt": prompt, "image": image_list, "answer": example["answer"]}


def apply_template(example):
    example["prompt"] = tokenizer.apply_chat_template(
        example["prompt"],
        tokenize=False,
        add_generation_prompt=False 
    )
    return example



dataset = dataset.map(make_conversation)
dataset = dataset.map(apply_template)

if (
        isinstance(images, (list, tuple))
        and all(isinstance(images_i, (list, tuple)) for images_i in images)
        and all(is_valid_list_of_images(images_i) for images_i in images)
    ):
        return images

    # If it's a list of images, it's a single batch, so convert it to a list of lists
    if isinstance(images, (list, tuple)) and is_valid_list_of_images(images):
        if is_pil_image(images[0]) or images[0].ndim == expected_ndims:
            return [images]
        if images[0].ndim == expected_ndims + 1:
            return [list(image) for image in images]

    # If it's a single image, convert it to a list of lists
    if is_valid_image(images):
        if is_pil_image(images) or images.ndim == expected_ndims:
            return [[images]]
        if images.ndim == expected_ndims + 1:
            return [list(images)]