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