from config import LinearMappingConfig from transformers import ( GPT2TokenizerFast, GPT2LMHeadModel, AutoModel, CLIPVisionModel, AutoProcessor, BatchEncoding, ) from transformers.models.gpt2.modeling_gpt2 import GPT2DoubleHeadsModelOutput import torch import torch.nn as nn from typing import List, Optional, Union, Tuple, Dict from torchvision.transforms import CenterCrop, ConvertImageDtype, Normalize, Resize from torchvision.transforms.functional import InterpolationMode class Transform(torch.nn.Module): def __init__(self, image_size, mean, std): super().__init__() self.transforms = torch.nn.Sequential( Resize([image_size], interpolation=InterpolationMode.BICUBIC, antialias=True), CenterCrop(image_size), ConvertImageDtype(torch.float32), Normalize(mean, std), ) def forward(self, x) -> torch.Tensor: """`x` should be an instance of `PIL.Image.Image`""" with torch.no_grad(): x = self.transforms(x) return x class LinearMappingProcessor: """ A combination of ImageProcessor and GPT2TokenizerFast """ def __init__(self, config: LinearMappingConfig): self.image_processor = AutoProcessor.from_pretrained(config.image_model) self.tokenizer = GPT2TokenizerFast.from_pretrained("gpt2") self.add_image_token = config.add_image_token if config.add_image_token: self.tokenizer.add_special_tokens({"cls_token": "||"}) self.tokenizer.pad_token_id = self.tokenizer.eos_token_id self.tokenizer.padding_side = "right" self.prefix_length = config.prefix_length def __call__(self, texts=None, images=None, return_tensors="pt", **kwargs): """ The processor assumes that images and texts are of the same number """ if len(texts) == 0: # empty strings should be None texts = None if images is not None: image_features = self.image_processor(images=images, return_tensors=return_tensors, **kwargs) image_features["attention_mask"] = torch.ones(image_features.pixel_values.size(0), self.prefix_length).to(dtype=torch.int64) if texts is None and self.add_image_token: texts = [self.tokenizer.cls_token for _ in range(image_features.pixel_values.size(0))] elif texts is not None and self.add_image_token: if isinstance(texts, str): texts = [texts] texts = [self.tokenizer.cls_token + text for text in texts] elif texts is None: texts = self.tokenizer.bos_token if texts is not None: encoding = self.tokenizer(texts, return_tensors=return_tensors, **kwargs) if texts is not None and images is not None: encoding["pixel_values"] = image_features.pixel_values encoding["attention_mask"] = torch.cat([ image_features["attention_mask"], encoding["attention_mask"] ], dim=1).to(dtype=torch.long) # create attention mask for images return encoding elif texts is not None: return encoding else: return BatchEncoding(data=dict(**image_features), tensor_type=return_tensors) def batch_decode(self, *args, **kwargs): """ This method forwards all its arguments to GPT2TokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please refer to the docstring of this method for more information. """ return self.tokenizer.batch_decode(*args, **kwargs) def decode(self, *args, **kwargs): """ This method forwards all its arguments to GPT2TokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to the docstring of this method for more information. """ return self.tokenizer.decode(*args, **kwargs) class ImagePrefix(nn.Module): """ Converts pixel values to prefix image prompts that are later fed to a LLM """ def __init__(self, config: LinearMappingConfig): super().__init__() self.encoder = AutoModel.from_pretrained(config.image_model) if "clip" in config.image_model: self.encoder = CLIPVisionModel.from_pretrained(config.image_model) if config.freeze_image_model: for param in self.encoder.parameters(): param.requires_grad = False self.linear = nn.Linear(config.image_hidden_size, config.text_hidden_size) self.ln = nn.LayerNorm(config.text_hidden_size) def forward( self, pixel_values: torch.Tensor # B x C x H x W ) -> torch.Tensor: prefixes = self.encoder(pixel_values).last_hidden_state # B x N x D prefix_prompts = self.linear(prefixes) return self.ln(prefix_prompts) class LinearMapping(nn.Module): def __init__(self, config: LinearMappingConfig): super().__init__() self.image_prefix = ImagePrefix(config) self.language_model = GPT2LMHeadModel.from_pretrained(config.text_model) self.processor = LinearMappingProcessor(config) self.tokenizer = self.processor.tokenizer self.image_processor = self.processor.image_processor self.add_image_token = config.add_image_token if config.add_image_token: self.language_model.resize_token_embeddings(len(self.tokenizer)) if config.freeze_text_model: for module in self.language_model.modules(): if not isinstance(module, nn.LayerNorm) or config.freeze_ln: for param in module.parameters(): param.requires_grad = False if config.add_image_token: # create a gradient mask for the lm_head weight and bias and hook it self.language_model.lm_head.weight.requires_grad = True self.weight_gradient_mask = nn.Parameter(torch.zeros_like(self.language_model.lm_head.weight), requires_grad=False) self.weight_gradient_mask[-1, :] = 1.0 self.language_model.lm_head.weight.register_hook(lambda grad: grad.mul_(self.weight_gradient_mask)) def prepare_text_inputs(self, input_ids: torch.Tensor) -> torch.Tensor: return self.language_model.transformer.wte(input_ids.to(dtype=torch.int64)) def prepare_inputs( self, input_ids: Optional[torch.Tensor], pixel_values: Optional[torch.Tensor] ) -> Dict: """ Prepare captions and pixel values for training. It takes the captions' input ids and turn them into input embeddings and turns pixel values into prefix prompts. Then it concatenates them into one whole prompt batch. """ if input_ids is not None and pixel_values is not None: text_embeddings = self.prepare_text_inputs(input_ids) # B x T x D prefix_prompts = self.image_prefix(pixel_values) # B x V x D inputs_embeddings = torch.cat([prefix_prompts, text_embeddings], dim=1) prefix_labels = torch.zeros(prefix_prompts.shape[:2], device=prefix_prompts.device) - 100 labels = torch.cat([prefix_labels, input_ids], dim=1) # B x (V + T) for label in labels: for k, token in enumerate(label): if token == self.tokenizer.eos_token_id: label[k + 1:] = -100 break return {"hidden_states": inputs_embeddings, "labels": labels.to(dtype=torch.int64)} elif pixel_values is not None: prefix_prompts = self.image_prefix(pixel_values) # B x V x D prefix_labels = torch.zeros(prefix_prompts.shape[:2], device=prefix_prompts.device) - 100 return {"hidden_states": prefix_prompts, "labels": prefix_labels.to(dtype=torch.int64)} elif input_ids is not None: text_embeddings = self.prepare_text_inputs(input_ids) labels = input_ids.clone() for label in labels: for k, token in enumerate(label): if token == self.tokenizer.eos_token_id: label[k + 1:] = -100 break return {"hidden_states": text_embeddings, "labels": labels.to(dtype=torch.int64)} else: return {"hidden_states": None, "labels": None} @torch.no_grad() def generate( self, input_ids: Optional[torch.Tensor] = None, pixel_values: Optional[torch.Tensor] = None, **kwargs ): if pixel_values is None: return self.language_model.generate( input_ids=input_ids, **kwargs ) batch_size = pixel_values.size(0) past_input_ids = None if input_ids is None: if self.add_image_token: input_ids = torch.tensor([self.tokenizer.cls_token_id for _ in range(batch_size)]).view(batch_size, -1) else: input_ids = torch.tensor([self.tokenizer.bos_token_id for _ in range(batch_size)]).view(batch_size, -1) if input_ids.size(-1) <= 1: first_forward_outputs = self.forward( pixel_values=pixel_values ) else: first_forward_outputs = self.forward( pixel_values=pixel_values, input_ids=input_ids[:, :-1] ) past_input_ids = input_ids[:, :-1] input_ids = input_ids[:, -1].view(batch_size, -1) past_key_values = first_forward_outputs.past_key_values if kwargs.get("attention_mask", None) is None: attention_mask_size = (past_key_values[0][0].size(0), past_key_values[0][0].size(-2)) attention_mask = torch.ones(attention_mask_size, dtype=torch.int64) else: attention_mask = kwargs.pop("attention_mask") generated_token_ids = self.language_model.generate( past_key_values=past_key_values, input_ids=input_ids, attention_mask=attention_mask, **kwargs ) if past_input_ids is not None: generated_token_ids = torch.cat([past_input_ids, generated_token_ids], dim=-1) return generated_token_ids def forward( self, input_ids: Optional[torch.Tensor] = None, pixel_values: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, output_hidden_states: bool = True, output_attentions: bool = True, attention_mask: Optional[torch.Tensor] = None, return_dict: Optional[bool] = True, **kwargs ) -> Union[GPT2DoubleHeadsModelOutput, Tuple]: if (pixel_values is None and input_ids is None) and inputs_embeds is None: raise ValueError("You have to specify inputs") if inputs_embeds is not None and (pixel_values is not None or input_ids is not None): raise ValueError("Either inputs_embeds or (pixel_values and input_ids) should be specified, not both") inputs = self.prepare_inputs(input_ids, pixel_values) hidden_states = inputs.get('hidden_states', None) if inputs_embeds is None else inputs_embeds labels = inputs.get('labels', None) if labels is None else labels return self.language_model( inputs_embeds=hidden_states, labels=labels, output_hidden_states=output_hidden_states, output_attentions=output_attentions, attention_mask=attention_mask, return_dict=return_dict, **kwargs )