File size: 3,139 Bytes
eacbbc9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
"""
Model implementation for VQA
"""
import os
import json
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel, AutoConfig, ViTImageProcessor, ViTModel
class VQAModel(nn.Module):
"""Vision-Language model for Visual Question Answering"""
def __init__(self, config, num_answers):
super(VQAModel, self).__init__()
self.config = config
self.num_answers = num_answers
# Vision encoder
self.vision_config = AutoConfig.from_pretrained(config['vision_model'])
self.vision_encoder = ViTModel.from_pretrained(config['vision_model'])
# Text encoder
self.text_config = AutoConfig.from_pretrained(config['text_model'])
self.text_encoder = AutoModel.from_pretrained(config['text_model'])
# Projection layers
self.vision_projection = nn.Linear(
self.vision_config.hidden_size, config['hidden_size']
)
self.text_projection = nn.Linear(
self.text_config.hidden_size, config['hidden_size']
)
# Multimodal fusion
self.fusion = nn.Sequential(
nn.Linear(2 * config['hidden_size'], config['hidden_size']),
nn.LayerNorm(config['hidden_size']),
nn.GELU(),
nn.Dropout(config['dropout'])
)
# Answer prediction
self.classifier = nn.Sequential(
nn.Linear(config['hidden_size'], config['hidden_size']),
nn.LayerNorm(config['hidden_size']),
nn.GELU(),
nn.Dropout(config['dropout']),
nn.Linear(config['hidden_size'], num_answers)
)
# Answerable prediction
self.answerable_classifier = nn.Sequential(
nn.Linear(config['hidden_size'], config['hidden_size'] // 2),
nn.LayerNorm(config['hidden_size'] // 2),
nn.GELU(),
nn.Dropout(config['dropout']),
nn.Linear(config['hidden_size'] // 2, 2) # Binary classification
)
def forward(self, image_encodings, question_encodings):
"""Forward pass of the model"""
# Process image
vision_outputs = self.vision_encoder(**image_encodings)
vision_embeds = vision_outputs.last_hidden_state[:, 0] # CLS token
vision_embeds = self.vision_projection(vision_embeds)
# Process text
text_outputs = self.text_encoder(**question_encodings)
text_embeds = text_outputs.last_hidden_state[:, 0] # CLS token
text_embeds = self.text_projection(text_embeds)
# Combine modalities
multimodal_features = torch.cat([vision_embeds, text_embeds], dim=1)
fused_features = self.fusion(multimodal_features)
# Predict answers and answerable
answer_logits = self.classifier(fused_features)
answerable_logits = self.answerable_classifier(fused_features)
return {
'answer_logits': answer_logits,
'answerable_logits': answerable_logits,
'fused_features': fused_features
} |