dataset.py

from transformers import AutoProcessor
from PIL import Image
import os
import torch
import pickle

## ACTUAL INPUT CONSTRUCTION
BASE_SPEAKER_LEN = 787

def joint_listener_input(processor, context_images, description, device):
    # Preliminaries
    img_dir = "tangram_pngs"
    raw_images = process_images(img_dir, context_images)
    target_anno = description.lower()
    prompt = construct_listener_full_prompt(
        processor, target_anno, 0, "verbose_instruction"
    )

    # Listener processing
    outputs = processor(
        text=[prompt],
        images=[raw_images],
        return_tensors="pt"
    ).to(device)
    l_input_tokens = outputs['input_ids'][:, :-2]
    l_attn_mask = outputs['attention_mask'][:, :-2]
    l_attn_mask[(l_input_tokens == 0).bool()] = 0
    images = outputs['pixel_values']
    l_image_attn_mask = outputs['pixel_attention_mask']

    # Speaker processing
    prompts = []
    for i in range(10):
        prompt = construct_speaker_full_prompt(processor, description, i, "information_after")
        prompts.append(prompt)
    outputs = processor(
        text=prompts,
        images=[raw_images]*10,
        padding='longest',
        return_tensors="pt"
    ).to(device)

    s_input_tokens = outputs['input_ids'][:, :-1]
    s_attn_mask = outputs['attention_mask'][:, :-1]
    s_attn_mask[(s_input_tokens == 0).bool()] = 0
    s_image_attn_mask = outputs['pixel_attention_mask']
    s_target_tokens = outputs['input_ids'][:, 1:]
    s_target_mask = []
    for i in range(10):
        curr_mask = create_speaker_caption_mask(outputs['input_ids'][i], s_attn_mask[i])
        s_target_mask.append(curr_mask)
    s_target_mask = torch.stack(s_target_mask, dim=0)

    return images, l_input_tokens, l_attn_mask, l_image_attn_mask, s_input_tokens.unsqueeze(0), \
        s_attn_mask.unsqueeze(0), s_image_attn_mask.unsqueeze(0), s_target_mask.unsqueeze(0), \
        s_target_tokens.unsqueeze(0)

def joint_speaker_input(processor, image_paths, target_path, device):
    # Get the prompt
    img_dir = "tangram_pngs"
    raw_images = process_images(img_dir, image_paths)
    target_idx = image_paths.index(target_path)
    base_prompt = construct_speaker_base_prompt(processor, target_idx, "information_after", process=True)

    # Create the basic input
    outputs = processor(
        text=[base_prompt],
        images=[raw_images],
        return_tensors="pt"
    ).to(device)

    input_tokens = outputs['input_ids']
    attn_mask = outputs['attention_mask']
    attn_mask[(input_tokens == 0).bool()] = 0
    images = outputs['pixel_values']
    image_attn_mask = outputs['pixel_attention_mask']

    return input_tokens, attn_mask, images, image_attn_mask, torch.LongTensor([target_idx]).to(device)


## UTILITIES

def get_processor():
    checkpoint = "HuggingFaceM4/idefics2-8b"
    processor = AutoProcessor.from_pretrained(checkpoint, do_image_splitting=False,
                                               size={"longest_edge": 448, "shortest_edge": 224})
    return processor

def get_index_to_token():
    index_to_token_path = "index_to_token.pkl"
    with open(index_to_token_path, 'rb') as f:
        index_to_token = pickle.load(f)
    return index_to_token

def process_images(img_dir, context_images):
    raw_images = []
    for img in context_images:
        image_path = os.path.join(img_dir, img)
        raw_image = Image.open(image_path).convert('RGB')
        raw_images.append(raw_image)
    return raw_images

def create_speaker_caption_mask(all_token_ids, text_mask):
    # Overall token comp: pad + base + caption
    padding_tokens = torch.sum(all_token_ids == 0).item()
    caption_tokens = all_token_ids.shape[0] - (padding_tokens + BASE_SPEAKER_LEN)

    # Construct a mask where the last caption tokens are 1
    target_mask = torch.zeros_like(text_mask)
    target_mask[-caption_tokens:] = 1

    return target_mask.bool()

def construct_listener_full_prompt(processor, target_anno, target_idx, comprehension_prompt_type="verbose_instruction"):
    target_anno = target_anno.lower().strip()
    messages = []

    if comprehension_prompt_type == "verbose_instruction":
        # User side: Intro
        messages.append(
            {
                "role" : "user",
                "content" : [
                    {"type" : "text", "text" : "You will be presented with a sequence of 10 images and a caption describing exactly one of them. "},
                    {"type" : "text", "text" : "Your task is to guess which image the caption describes. "},
                ]
            }
        )

        # User side: Images
        for i in range(10):
            if i == 0:
                messages[0]["content"].append({"type" : "text", "text" : f" Image {i}: "})
            else:
                messages[0]["content"].append({"type" : "text", "text" : f", Image {i}: "})
            messages[0]["content"].append({"type" : "image"})

        # User side: Caption
        messages[0]["content"].append({"type" : "text", "text" : f". Caption: {target_anno}"})
        messages[0]["content"].append({"type" : "text", "text" : f" Does this caption describe Image 0, 1, 2, 3, 4, 5, 6, 7, 8 or 9?"})

        # Model side: Guess
        messages.append(
            {
                "role" : "assistant",
                "content" : [
                    {"type" : "text", "text" : f"The caption describes Image {target_idx}"}
                ]
            }
        )
    else:
        assert(False)

    return processor.apply_chat_template(messages, add_generation_prompt=False).strip() 

def construct_speaker_full_prompt(processor, target_anno, target_idx,
                                  generation_prompt_type="information_after"):
    messages = construct_speaker_base_prompt(processor, target_idx, generation_prompt_type)

    # Assistant response
    target_anno = target_anno.lower().strip()
    messages.append(
        {
            "role" : "assistant",
            "content" : [
                {"type" : "text", "text" : target_anno}
            ]
        }
    )

    return processor.apply_chat_template(messages, add_generation_prompt=False).strip()

def construct_speaker_base_prompt(processor, target_idx, generation_prompt_type="information_after", process=False):
    messages = []

    if generation_prompt_type == "information_after":
        # User side: Intro
        messages.append(
            {
                "role" : "user",
                "content" : [
                    {"type" : "text", "text" : "You will be presented with a sequence of 10 images and be assigned a target image. "},
                    {"type" : "text", "text" : "Your task is to produce a caption for your target image such that anyone could guess the image from your description. "},
                ]
            }
        )

        # User side: Images
        for i in range(10):
            if i == 0:
                messages[0]["content"].append({"type" : "text", "text" : f" Image {i}: "})
            else:
                messages[0]["content"].append({"type" : "text", "text" : f", Image {i}: "})
            messages[0]["content"].append({"type" : "image"})

        # User side: Target assignment
        messages[0]["content"].append({"type" : "text", "text" : f". Your target image is Image {target_idx}. Produce your caption now."})
    else:
        assert(False)

    if process:
        prompt = processor.apply_chat_template(messages, add_generation_prompt=True).strip()
        return prompt
    else:
        return messages

def process_idefics_listener_generation_input(speaker_context, captions, processor, img_dir, num_samples, device):
    # First construct the prompts
    prompts, raw_images = get_listener_generation_prompts(speaker_context, captions, num_samples, img_dir, processor)

    # Process the prompts
    listener_inputs = processor(
        text=prompts,
        images=raw_images,
        padding='longest',
        return_tensors='pt'
    )

    input_tokens = listener_inputs['input_ids'][:, :-2].to(device)
    attn_mask = listener_inputs['attention_mask'][:, :-2].to(device)
    attn_mask[input_tokens == 0] = 0
    images = listener_inputs['pixel_values'].to(device)
    image_attn_mask = listener_inputs['pixel_attention_mask'].to(device)

    return input_tokens, attn_mask, images, image_attn_mask

def get_listener_generation_prompts(speaker_contexts, captions, num_samples, img_dir, processor):
    prompts = []
    all_raw_images = []

    for i, speaker_context in enumerate(speaker_contexts):
        raw_images = process_images(img_dir, speaker_context)
        for j in range(num_samples):
            curr_idx = i * num_samples + j
            caption = captions[curr_idx]
            prompt = construct_listener_full_prompt(processor, caption, 0, "verbose_instruction")

            prompts.append(prompt)
            all_raw_images.append(raw_images)
    return prompts, all_raw_images