vision/m4/models/vgpt2/evaluation_image_caption_matching_vgpt2.py

import os
from itertools import chain
from typing import Dict, List, Optional

import torch
from torch.nn import CrossEntropyLoss
from transformers import AutoTokenizer

from m4.evaluation.custom_metrics.image_caption_matching_metrics import MetricsImageCaptionMatching
from m4.evaluation.tasks import BaseTaskImageCaptionMatching, Predictor
from m4.evaluation.utils import EvaluationVersion
from m4.training.utils import (
    FAKE_TOKEN_AROUND_IMAGE_V1,
    FAKE_TOKEN_AROUND_IMAGE_V2,
    IMAGE_TOKEN,
    build_image_transform,
)


class Vgpt2ImageCaptionMatching(BaseTaskImageCaptionMatching):
    model_class: str = "VGPT2LMHeadModel"
    predictor_class: Predictor = Predictor.in_contexter
    target_keys: List[str] = ["example_ids", "caption_ids", "image_ids"]
    buckets_keys: List[str] = []
    # Buckets are optionally populated for classification in context. They are only useful when it is useful to get results bucket. A bucket is typically a certain slice of the dataset (for instance, all instances where age=30).
    mapping_class_names_to_prompt_names: Optional[Dict[str, str]] = None
    prompt_templates_dict: Dict[int, Dict[str, str]] = {}
    mapping_class_prompt_name_id_to_prompt_template_id: Optional[Dict[int, int]] = None
    tokenizer_max_seq_len = 1024

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
        self.tokenizer_name = kwargs.pop("tokenizer_name")
        evaluation_version = kwargs.pop("evaluation_version")
        tokenizer_use_fast = kwargs.pop("tokenizer_use_fast", False)
        image_size = kwargs.pop("image_size")
        vision_encoder_type = kwargs.pop("vision_encoder_type")
        self.image_seq_len = kwargs.pop("image_seq_len")
        self.image_transform = build_image_transform(
            max_image_size=image_size, image_size=None, eval=True, vision_encoder_type=vision_encoder_type
        )

        self.tokenizer = AutoTokenizer.from_pretrained(
            self.tokenizer_name, truncation_side="left", use_fast=tokenizer_use_fast, token=os.getenv("HF_TOKEN", True)
        )
        self.tokenizer.padding_side = "left"
        self.image_token = IMAGE_TOKEN
        if evaluation_version == EvaluationVersion.v1:
            self.token_around_image = FAKE_TOKEN_AROUND_IMAGE_V1
        elif evaluation_version == EvaluationVersion.v2:
            self.token_around_image = FAKE_TOKEN_AROUND_IMAGE_V2
        else:
            raise ValueError(f"Invalid evaluation version: {evaluation_version}")

        nb_captions = len(self.caption_column_names)
        nb_images = len(self.image_column_names)
        self.captions_images_order_per_ex = [
            (caption_idx, image_idx) for caption_idx in range(nb_captions) for image_idx in range(nb_images)
        ]

    def get_info_from_dataset(self, dataset):
        pass

    def get_data_collator(self, **kwargs):
        def data_collator(batch):
            exs = {key: [ex[key] for ex in batch] for key in batch[0].keys()}
            batch = self.prepare_dataset(exs, **kwargs)
            return batch

        return data_collator

    def _split_array(self, array, nb_combinations):
        total_elements = len(array)
        elements_per_combination = total_elements // nb_combinations

        splitted_array = [
            array[i : i + elements_per_combination] for i in range(0, total_elements, elements_per_combination)
        ]

        return splitted_array

    def prepare_dataset(self, exs: Dict, **kwargs) -> Dict:
        """
        Prepare batch of examples.

        """

        prompt_template_id: int = kwargs["prompt_template_id"]

        nb_exs = len(exs["id"])
        nb_captions = len(self.caption_column_names)
        nb_images = len(self.image_column_names)

        # If we have caption_column_names = ["caption_0", "caption_1"] and image_column_names= ["image_0", "image_1"]. We get the sequence [caption_0, caption_0, caption_1, caption_1]
        general_dict = {"tested_prompts": [], "caption_ids": [], "image_ids": [], "ex_ids": []}
        for idx_ex in range(nb_exs):
            for caption_idx, caption_column in enumerate(self.caption_column_names):
                for image_idx in range(nb_images):
                    tested_prompt = self._create_example_prompt(
                        prompt_template_id=prompt_template_id,
                        caption=exs[caption_column][idx_ex],
                    )
                    general_dict["tested_prompts"].append(tested_prompt)
                    general_dict["caption_ids"].append(caption_idx)
                    general_dict["image_ids"].append(image_idx)
                    general_dict["ex_ids"].append(exs["id"][idx_ex])

        tot_texts = [
            self._create_prefix_prompt(prompt_template_id=prompt_template_id) + tested_prompt
            for tested_prompt in general_dict["tested_prompts"]
        ]

        tot_texts = [text.strip() for text in tot_texts]
        # Tokenize and masks
        tokens = self.tokenizer(
            tot_texts,
            return_tensors="pt",
            truncation=True,
            max_length=self.tokenizer_max_seq_len,
            padding=True,
            add_special_tokens=False,
        )

        general_dict["input_ids"] = [tokens.input_ids[idx] for idx in range(len(tot_texts))]
        general_dict["attention_mask"] = [tokens.attention_mask[idx] for idx in range(len(tot_texts))]

        # If we have caption_column_names = ["caption_0", "caption_1"] and image_column_names= ["image_0", "image_1"]. We get the sequence image_0, image_1, image_0, image_1
        pixel_values_dict = {"pixel_values": [], "caption_ids": [], "image_ids": [], "ex_ids": []}
        for idx_ex in range(nb_exs):
            for caption_idx in range(nb_captions):
                for image_idx, col in enumerate(self.image_column_names):
                    pixel_values_dict["pixel_values"].append(self.image_transform(exs[col][idx_ex]).unsqueeze(0))
                    pixel_values_dict["caption_ids"].append(caption_idx)
                    pixel_values_dict["image_ids"].append(image_idx)
                    pixel_values_dict["ex_ids"].append(exs["id"][idx_ex])

        # ---- Sanity check ----
        assert pixel_values_dict["ex_ids"] == general_dict["ex_ids"]
        nb_combinations = nb_captions * nb_images
        sample_pixel_captions_ids = pixel_values_dict["caption_ids"][:nb_combinations]
        sample_pixel_image_ids = pixel_values_dict["image_ids"][:nb_combinations]
        sample_general_captions_ids = general_dict["caption_ids"][:nb_combinations]
        sample_general_image_ids = general_dict["image_ids"][:nb_combinations]
        self.captions_images_order_per_ex
        for idx in range(nb_combinations):
            expected_caption_idx, expected_image_idx = self.captions_images_order_per_ex[idx]
            assert sample_pixel_captions_ids[idx] == expected_caption_idx
            assert sample_general_captions_ids[idx] == expected_caption_idx
            assert sample_pixel_image_ids[idx] == expected_image_idx
            assert sample_general_image_ids[idx] == expected_image_idx
        # ---- Sanity check ----

        general_dict["ex_ids"] = self._split_array(general_dict["ex_ids"], nb_exs)
        general_dict["caption_ids"] = self._split_array(general_dict["caption_ids"], nb_exs)
        general_dict["image_ids"] = self._split_array(general_dict["image_ids"], nb_exs)
        general_dict["input_ids"] = self._split_array(general_dict["input_ids"], nb_exs)
        pixel_values_dict["pixel_values"] = self._split_array(pixel_values_dict["pixel_values"], nb_exs)
        general_dict["attention_mask"] = self._split_array(general_dict["attention_mask"], nb_exs)

        return {
            "example_ids": general_dict["ex_ids"],
            "caption_ids": general_dict["caption_ids"],
            "image_ids": general_dict["image_ids"],
            "input_ids": general_dict["input_ids"],
            "attention_mask": general_dict["attention_mask"],
            "pixel_values": pixel_values_dict["pixel_values"],
        }

    def _create_example_prompt(self, prompt_template_id, caption):
        if self.bool_instruct_templates:
            prompt_templates_dict = self.prompt_templates_dict_instruct
        else:
            prompt_templates_dict = self.prompt_templates_dict
        prompt_template = prompt_templates_dict[prompt_template_id]["example"]

        template_kwargs = {
            "bos_token": self.tokenizer.bos_token,
            "eos_token": self.tokenizer.eos_token,
            "image_token": self.image_token * self.image_seq_len,
            "token_around_image": self.token_around_image,
            "caption": caption,
        }

        prompt = prompt_template.format(**template_kwargs)
        return prompt

    def _create_prefix_prompt(self, prompt_template_id):
        if self.bool_instruct_templates:
            prompt_templates_dict = self.prompt_templates_dict_instruct
        else:
            prompt_templates_dict = self.prompt_templates_dict
        prompt_template = prompt_templates_dict[prompt_template_id]["prefix"]
        if prompt_template is None:
            return ""
        else:
            prompt = prompt_template.format(
                bos_token=self.tokenizer.bos_token,
                eos_token=self.tokenizer.eos_token,
            )
        return prompt

    def predict(self, **kwargs):
        model = kwargs["model"]
        input_ids = torch.stack(list(chain.from_iterable(kwargs["input_ids"]))).to(model.device)
        attention_mask = torch.stack(list(chain.from_iterable(kwargs["attention_mask"]))).to(model.device)

        pv = list(chain.from_iterable(kwargs["pixel_values"]))
        total_batch_size = len(pv)
        max_num_images = max([i.size(0) for i in pv])
        max_height = max([i.size(2) for i in pv])
        max_width = max([i.size(3) for i in pv])
        pixel_values = torch.zeros(total_batch_size, max_num_images, 3, max_height, max_width)
        pixel_attention_mask = torch.zeros(total_batch_size, max_num_images, max_height, max_width, dtype=torch.bool)
        for idx, sample_images in enumerate(pv):
            im_batch_height, im_batch_width = sample_images.size()[2:]
            pixel_values[idx, :, :, :im_batch_height, :im_batch_width] = sample_images
            pixel_attention_mask[idx, :, :im_batch_height, :im_batch_width] = True
        pixel_values = pixel_values.to(model.device)
        pixel_attention_mask = pixel_attention_mask.to(model.device)

        with torch.no_grad():
            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                pixel_values=pixel_values,
                pixel_attention_mask=pixel_attention_mask,
            )
        outputs.input_ids = input_ids
        outputs.attention_mask = attention_mask
        return outputs

    def format_model_outputs_to_predictions(self, outputs) -> torch.Tensor:
        batch_size = outputs.logits.shape[0]
        # Shift so that tokens < n predict n
        shift_logits = outputs.logits[..., :-1, :].contiguous()
        shift_labels = outputs.input_ids[..., 1:].contiguous()
        shift_attention_mask = outputs.attention_mask[:, 1:].contiguous()
        # Flatten the tokens
        loss_fct = CrossEntropyLoss(reduction="none")
        log_probs = -loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
        log_probs = log_probs.view(batch_size, -1)
        masked_log_probs = shift_attention_mask * log_probs
        score_per_example = masked_log_probs.sum(dim=-1)
        if self.length_normalize:
            score_per_example = score_per_example / shift_attention_mask.sum(dim=-1)

        nb_combinations = len(self.image_column_names) * len(self.caption_column_names)
        nb_exs = batch_size // nb_combinations
        splitted_scores_per_example = self._split_array(score_per_example.tolist(), nb_exs)

        return splitted_scores_per_example

    def add_batch_metric(self, metric, **kwargs):
        outputs = self.predict(**kwargs)
        splitted_scores_per_example = self.format_model_outputs_to_predictions(outputs)
        additional_args = {key: kwargs[key] for key in self.target_keys}
        metric.add_batch(
            splitted_scores_per_example=splitted_scores_per_example,
            **additional_args,
        )
        return metric


class WinogroundVgpt2ImageCaptionMatchingAccWithKLAndEntropy(Vgpt2ImageCaptionMatching):
    dataset_name: str = "facebook/winoground"
    metric_name: str = "ImageCaptionMatchingMetrics"
    metric_kwargs = {
        "metrics": [
            MetricsImageCaptionMatching.TEXT_SCORE,
            MetricsImageCaptionMatching.IMAGE_SCORE,
            MetricsImageCaptionMatching.GROUP_SCORE,
        ]
    }
    # support split names are never used for this dataset
    default_query_split_name: str = "test"
    default_support_split_name: str = "test"
    test_support_split_name: str = "test"
    image_column_names: List[str] = ["image_0", "image_1"]
    id_column_name: str = "id"
    caption_column_names: List[str] = ["caption_0", "caption_1"]
    length_normalize: bool = True
    prompt_templates_dict = {
        0: {
            "prefix": None,
            "example": "{token_around_image}{image_token}{token_around_image}{caption}",
        },
    }