| |
| """ |
| Tiny VQAScore Model Wrapper |
| """ |
|
|
| import torch |
| import torch.nn as nn |
| from PIL import Image |
| import numpy as np |
|
|
| class TinyVQAScore: |
| """A tiny random version of the VQAScore model.""" |
|
|
| def __init__(self, model="tiny-random", device="cpu"): |
| self.device = torch.device(device) |
| self.model = self._create_tiny_model() |
| self.model.to(self.device) |
| self.model.eval() |
|
|
| def _create_tiny_model(self): |
| class TinyCLIPT5(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.vision_encoder = nn.Sequential( |
| nn.Conv2d(3, 64, kernel_size=16, stride=16), |
| nn.AdaptiveAvgPool2d((1, 1)), |
| nn.Flatten(), |
| nn.Linear(64, 256) |
| ) |
| self.text_encoder = nn.Sequential( |
| nn.Embedding(32128, 256), |
| nn.LayerNorm(256), |
| nn.TransformerEncoderLayer(d_model=256, nhead=8, dim_feedforward=512, dropout=0.1, batch_first=True) |
| ) |
| self.multimodal_projector = nn.Sequential( |
| nn.Linear(256, 128), nn.GELU(), |
| nn.Linear(128, 64), nn.GELU(), |
| nn.Linear(64, 1) |
| ) |
| self._init_weights() |
|
|
| def _init_weights(self): |
| for module in self.modules(): |
| if isinstance(module, (nn.Linear, nn.Conv2d)): |
| nn.init.xavier_uniform_(module.weight, gain=0.1) |
| if module.bias is not None: |
| nn.init.uniform_(module.bias, -0.1, 0.1) |
| elif isinstance(module, nn.Embedding): |
| nn.init.uniform_(module.weight, -0.1, 0.1) |
|
|
| def forward(self, pixel_values, input_ids): |
| vision_features = self.vision_encoder(pixel_values) |
| text_features = self.text_encoder(input_ids) |
| text_features = text_features.mean(dim=1) |
| combined_features = vision_features + text_features |
| score = self.multimodal_projector(combined_features) |
| return score.squeeze(-1) |
|
|
| return TinyCLIPT5() |
|
|
| def score(self, image, question): |
| if isinstance(image, Image.Image): |
| image = image.resize((224, 224)) |
| image_tensor = torch.from_numpy(np.array(image)).float() |
| image_tensor = image_tensor.permute(2, 0, 1).unsqueeze(0) / 255.0 |
| else: |
| image_tensor = image |
|
|
| input_ids = torch.randint(0, 32128, (1, 10)).to(self.device) |
|
|
| with torch.no_grad(): |
| score = self.model(image_tensor.to(self.device), input_ids) |
|
|
| return torch.sigmoid(score).item() |
|
|
| if __name__ == "__main__": |
| |
| model = TinyVQAScore(device="cpu") |
| test_image = Image.new('RGB', (224, 224), color='red') |
| score = model.score(test_image, "What color is this image?") |
| print(f"Test score: {score}") |
|
|
