from datasets import load_dataset, features


def format(examples):
    """
    Convert prompt from "xxx" to [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "xxx"}]}]
    and chosen and rejected from "xxx" to [{"role": "assistant", "content": [{"type": "text", "text": "xxx"}]}].
    Images are wrapped in a list.
    """
    output = {"images": [], "prompt": [], "chosen": [], "rejected": []}
    for image, question, chosen, rejected in zip(examples["image"], examples["question"], examples["chosen"], examples["rejected"]):
        prompt = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": question}]}]
        chosen = [{"role": "assistant", "content": [{"type": "text", "text": chosen}]}]
        rejected = [{"role": "assistant", "content": [{"type": "text", "text": rejected}]}]
        output["images"].append([image])
        output["prompt"].append(prompt)
        output["chosen"].append(chosen)
        output["rejected"].append(rejected)
    return output


dataset = load_dataset("openbmb/RLAIF-V-Dataset", split="train")
cols = dataset.column_names
dataset = dataset.map(format, batched=True, writer_batch_size=4, batch_size=4, remove_columns=cols)
f = dataset.features
f["images"] = features.Sequence(features.Image(decode=True))  # to avoid bytes
dataset = dataset.cast(f)
dataset = dataset.train_test_split(test_size=0.05)
dataset.push_to_hub("HuggingFaceH4/rlaif-v_formatted")