Spaces:
Sleeping
Sleeping
| import torch | |
| from transformers import ViTModel, ViTFeatureExtractor, GPT2LMHeadModel, GPT2Tokenizer | |
| from config.config import Config | |
| from torchvision import transforms | |
| class ImageCaptioningModel: | |
| def __init__(self): | |
| """Initialize the ViT and GPT-2 models for image captioning.""" | |
| self.device = Config.DEVICE | |
| self.vit_model = ViTModel.from_pretrained(Config.VIT_MODEL).to(self.device) | |
| self.feature_extractor = ViTFeatureExtractor.from_pretrained(Config.VIT_MODEL) | |
| self.gpt2_model = GPT2LMHeadModel.from_pretrained(Config.GPT2_MODEL).to(self.device) | |
| self.tokenizer = GPT2Tokenizer.from_pretrained(Config.GPT2_MODEL) | |
| self.tokenizer.pad_token = self.tokenizer.eos_token | |
| def extract_image_features(self, images): | |
| """Extract features from images using ViT.""" | |
| pixel_values = self.feature_extractor(images=images, return_tensors="pt", do_rescale=False).pixel_values.to(self.device) | |
| with torch.no_grad(): | |
| outputs = self.vit_model(pixel_values) | |
| return outputs.last_hidden_state[:, 0, :] # [batch_size, hidden_size] | |
| def prepare_gpt2_inputs(self, image_features, captions): | |
| """Prepare GPT-2 inputs.""" | |
| # Tokenize the captions | |
| tokenized_captions = self.tokenizer(captions, padding="longest", truncation=True, | |
| max_length=Config.MAX_SEQ_LEN, return_tensors="pt").to(self.device) | |
| # Get the word embeddings for the tokens | |
| token_embeddings = self.gpt2_model.transformer.wte(tokenized_captions['input_ids']) | |
| # Concatenate image features with token embeddings | |
| image_features = image_features.unsqueeze(1) # Reshape to [batch_size, 1, hidden_size] | |
| inputs_embeds = torch.cat((image_features, token_embeddings), dim=1) # Concatenate along the sequence dimension | |
| # Adjust input_ids to account for the image feature token | |
| batch_size = image_features.shape[0] | |
| image_token_id = torch.full((batch_size, 1), fill_value=self.tokenizer.bos_token_id, device=self.device) | |
| input_ids = torch.cat((image_token_id, tokenized_captions['input_ids']), dim=1) | |
| # Adjust attention_mask to account for the image feature token | |
| image_attention = torch.ones((batch_size, 1), device=self.device) | |
| attention_mask = torch.cat((image_attention, tokenized_captions['attention_mask']), dim=1) | |
| return inputs_embeds, input_ids, attention_mask | |
| def save(self, path): | |
| """Save model to disk.""" | |
| self.gpt2_model.save_pretrained(path) | |
| def load(self, path): | |
| """Load model from disk.""" | |
| self.gpt2_model = GPT2LMHeadModel.from_pretrained(path).to(self.device) | |
| # return self.gpt2_model | |