data/datasets.py

from __future__ import annotations

import logging
from collections.abc import Iterator
from typing import Optional

import torch
from PIL import Image
from torch.utils.data import Dataset

class BaseDataset(Dataset):
    def __init__(
        self,
        dataset,
        tokenizer,
        image_processor,
        relevance_min_rating: int = 1,
        image_correspondence_min_rating: int = 1,
        visual_dependency_min_rating: int = 1,
        formatting_min_rating: int = 1,
    ) -> None:
        self.dataset = dataset
        self.tokenizer = tokenizer
        self.image_processor = image_processor
        self.relevance_min_rating = relevance_min_rating
        self.image_correspondence_min_rating = image_correspondence_min_rating
        self.visual_dependency_min_rating = visual_dependency_min_rating
        self.formatting_min_rating = formatting_min_rating
        self.prefix_len = self._get_prefix_len()

    def __len__(self) -> int:
        return len(self.dataset)

    def _get_prefix_len(self) -> int:
        """
        Purpose:
            Count the number of tokens in the prefix that gets prepended to each assistant turn. 
            The prefix is <|im_start|>assistant\n, which consists of ['<|im_start|>', 'ass', 'istant', '\n'] tokens
        
        Parameters:
            None

        Returns:
            Integer representing the number of tokens that the prefix <|im_start|>assistant\n consists of        
        """
        random_string_5_letters = "xzyvd"
        random_string_chat_templated = self.tokenizer.apply_chat_template([{"role": "assistant", "content": random_string_5_letters}], tokenize=False, add_special_tokens=False)
        random_string_location = random_string_chat_templated.find(random_string_5_letters)
        return len(self.tokenizer.encode(random_string_chat_templated[:random_string_location]))

    def _get_messages(self, item: dict, image_token_counts: list[int]) -> list[dict]:
        """
        Purpose:
            Given a sample (item), creates a list of dictionaries. Each dictionary is of format 
            {"role": "user", "content": ...} or {"role": "assistant", "content": ...}.
            
            Prepends image tokens to "content" of the first message. 
            The number of prepended image tokens is the sum [ image_tokens_in_image_1 + ... + image_tokens_in_image_N ]
        
        Parameters:
         * item (dict) : a dictionary representing a sample from the FineVision dataset

         * image_token_counts (list) : a list, where image_token_counts[i] is the number of image tokens
                                       image i of the sample was decomposed into

        Returns:
            A list of dictionaries, where each dictionary is of format
            {"role": "user", "content": ...} or {"role": "assistant", "content": ...}.
        """
        messages = []
        for index, text in enumerate(item['texts']):
            try:
                if item.get('relevance_ratings') is not None and item['relevance_ratings'][index] is not None and item['relevance_ratings'][index] < self.relevance_min_rating:
                    continue
                if item.get('image_correspondence_ratings') is not None and item['image_correspondence_ratings'][index] is not None and item['image_correspondence_ratings'][index] < self.image_correspondence_min_rating:
                    continue
                if item.get('visual_dependency_ratings') is not None and item['visual_dependency_ratings'][index] is not None and item['visual_dependency_ratings'][index] < self.visual_dependency_min_rating:
                    continue
                if item.get('formatting_ratings') is not None and item['formatting_ratings'][index] is not None and item['formatting_ratings'][index] < self.formatting_min_rating:
                    continue
            except Exception as e:
                logging.warning(f"Error processing item: {item}, index: {index}: {e}")

            messages.append({"role": "user", "content": text['user']})
            messages.append({"role": "assistant", "content": text['assistant']})

        if len(messages) == 0:
            return messages

        # Safety check to ensure no image tokens are present in the text before adding them.
        for msg in messages:
            if self.tokenizer.image_token in msg["content"]:
                logging.warning(f"Found and removed an image token in the {msg['role']} text before adding the image string.")
                msg["content"] = msg["content"].replace(self.tokenizer.image_token, "")

        if len(image_token_counts) > 0:
            # Prepend image tokens to the content of first user message
            image_string = self.tokenizer.image_token * sum(image_token_counts) 
            messages[0]["content"] = image_string + messages[0]["content"]

        return messages

    def _process_images(self, images: list[Image.Image]) -> tuple[list[torch.Tensor], list[torch.Tensor], list[int]]:
        """
        Purpose:
            Takes in a list of list of PIL images. Returns a list, where element i
            is a sequence of flattened model-size patches from image i, as well as a list
            where element i is the number of image tokens/model-size patches in image i.

        Parameters:
         * self

         * images (list) : a list of PIL images

        Returns:
         * processed_images (list) : processed_images[i] is a 2D tensor of shape (num_model_patches, flat_model_patch_size),
                                     containing flat patches extracted from image i

         * model_patch_positions_list (list) : model_patch_positions_list[i] is a 2D tensor of shape (num_image_tokens, 2).
                                               It contains original 2D positions of flattened patches in image i.
                                    
         * image_token_counts (list) : image_token_counts[i] is the number of image tokens/model-size patches
                                       extracted from images[i]
        """
        processed_images = []
        model_patch_positions_list = []
        image_token_counts = []
        for image in images:
            if isinstance(image, Image.Image):
                if image.mode != 'RGB':
                    image = image.convert('RGB')
                
                processed_image, model_patch_positions, num_image_tokens = self.image_processor(image)
                
                processed_images.append(processed_image)
                model_patch_positions_list.append(model_patch_positions)
                image_token_counts.append(num_image_tokens)
            else:
                raise ValueError(f"Error processing image: {image}")
        return processed_images, model_patch_positions_list, image_token_counts


    def _prepare_inputs_and_loss_mask(
        self, messages: list[dict]
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Purpose:
            Given a list of messages representing a conversation, returns tokenized conversation (as tensor of integers),
            loss mask (as tensor of bools), and attention mask (as tensor of bools).

        Parameters: 
         * self

         * messages (list) : a list representing a single sample. It contains dictionaries
                             of format {"role": "user", "content": ...} or {"role": "assistant", "content": ...}.
               
        Returns:
            torch.tensor(conv_ids["input_ids"]) - a tensor of integers, representing tokenized conversation

            torch.tensor(mask).to(torch.bool) - mask that's True for tokens that may contribute to loss and False for tokens that may not

            torch.tensor(conv_ids["attention_mask"]) - mask of int64 that's 1 for tokens that may be attended to and 0 for tokens that may not
        """
        # conv_ids is a dictionary with keys "input_ids" and "attention_mask"
        # conv_ids["input_ids"] is a list of integers, representing a tokenized conversation
        # conv_ids["attention_mask"] is a list of integers (only 0 or 1) s.t. conv_ids["attention_mask"][i] is 
        # 0 if other tokens may attend to token i, and 1 otherwise
        conv_ids = self.tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            add_special_tokens=False,
            return_dict=True,
        )
        # Initialize the mask to an array of 0s - initially probability distribution for NEITHER token 
        # may contribute to loss
        mask = [0] * len(conv_ids["input_ids"])

        # Locate each assistant turn and flip its mask to 1
        cursor = 0
        for msg in messages:
            # Get a list of token ids for a single message
            segment_ids = self.tokenizer.apply_chat_template(
                [msg], tokenize=True, add_special_tokens=False, return_dict=True
            )["input_ids"]
            # Count the number of token ids in that message
            seg_len = len(segment_ids)

            # Determine positions of tokens from assistant's response,
            # and set mask to 1 at those positions (i.e. specify that predictions for these tokens may contribute to loss)
            if msg["role"] == "assistant":
                start = cursor + self.prefix_len
                end   = cursor + seg_len
                mask[start:end] = [1] * (end - start)  # attend to these tokens

            cursor += seg_len
        
        return torch.tensor(conv_ids["input_ids"]), torch.tensor(mask).to(torch.bool), torch.tensor(conv_ids["attention_mask"])


class VQADataset(BaseDataset):  # Visual Question Answering Dataset
    def iter_for_worker(self) -> Iterator[Optional[dict]]:  # with iterable datasets, each worker gets different shards
        for data in self.dataset:
            yield self._process_data(data)

    def __getitem__(self, idx: int) -> Optional[dict]:
        item = self.dataset[idx]
        return self._process_data(item)

    def _process_data(self, item: dict) -> Optional[dict]:
        """
        Purpose:
            Process a single sample from a dataset.

        Parameters:
         * item (dict) :  dictionary with keys "images" and "texts". 

        Returns:
            A dictionary with keys 'images', 'input_ids', 'attention_mask', 'labels'

            * out['images'] (list) : a list of tensors, where tensor i contains flattened
                                     patch embeddings from image i

            * out['model_patch_positions'] (list) : a list of tensors, where tensor i is a 2D tensor of shape (num_image_tokens, 2).
                                                    It contains original 2D positions of flattened patches in image i. 

            * out['input_ids'] (torch.Tensor) : a tensor of integers, representing the tokenized conversation.
                                                Includes image tokens.

            * out['attention_mask'] (torch.Tensor) : a tensor of bools. out['attention_mask'] is 1 if token i may be attended to
                                                     and 0 otherwise

            * out['labels'] (torch.Tensor) : a tensor of labels for calculating loss 
        """
        # Handle images (should be a list)
        if item['images'] is None:
            images_data = []
        else:
            images_data = item['images']
            if not isinstance(images_data, list):
                images_data = [images_data]

        processed_images = []
        model_patch_positions_list = []
        image_token_counts = []
        if images_data: # Only process if there are images
            processed_images, model_patch_positions_list, image_token_counts = self._process_images(images_data)

        messages = self._get_messages(item, image_token_counts)

        if len(messages) == 0:
            return None

        input_ids, mask, attention_mask = self._prepare_inputs_and_loss_mask(messages)
        labels = self._get_labels(input_ids, mask)

        return {
            "images": processed_images,
            "input_ids": input_ids,
            "model_patch_positions": model_patch_positions_list,
            "attention_mask": attention_mask,
            "labels": labels,
        }

    def _get_labels(self, input_ids: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        """
        Purpose:
            Given a tensor of token ids, return a tensor where element i is
            element (i+1) from the input tensor, or -100 if the prediction about token (i+1)
            (made by token i) may not contribute to loss.

        Parameters:
         * input_ids (torch.Tensor) - a tensor of integers, representing a sequence of token ids

         * mask (torch.Tensor) - a tensor of bools. mask[i] is True if prediction about token i may contribute
                                 to loss and 0 otherwise

        Returns:
            A tensor of integers described above.                      
        """
        labels = input_ids.clone().masked_fill(~mask, -100)
        labels = labels.roll(-1) # Shift labels for causal LM
        labels[-1] = -100 # Last token has no target
        
        return labels