Create VQA.py

VQA.py

@@ -0,0 +1,56 @@
+from transformers import AutoModel
+import torch
+import torch.nn as nn
+from typing import Optional
+fin= open("answer_space.txt")
+answer_space= fin.read().splitlines()
+
+class VQA(nn.Module):
+    def __init__(self, 
+                 text_encoder_name= 'bert-base-uncased', 
+                 image_encoder_name= 'google/vit-base-patch16-224-in21k',
+                 num_labels= 582):
+        super(VQA, self).__init__()
+        self.num_labels= num_labels
+        self.text_encoder_name= text_encoder_name
+        self.image_encoder_name= image_encoder_name
+        self.text_encoder= AutoModel.from_pretrained(self.text_encoder_name)
+        self.image_encoder= AutoModel.from_pretrained(self.image_encoder_name)
+    
+        # 768 + 768
+        self.combine= nn.Sequential(
+            nn.Linear(self.text_encoder.config.hidden_size+ self.image_encoder.config.hidden_size, 1059),
+            nn.ReLU(),
+            nn.Dropout(0.6)
+        )
+        self.layer1= nn.Linear(1059, 1059)
+        self.classifier= nn.Linear(1059, self.num_labels)
+        self.criterion= nn.CrossEntropyLoss()
+
+    def forward(self, input_ids, pixel_values, attention_mask, token_type_ids, labels: Optional[torch.LongTensor]= None):
+        encoded_text= self.text_encoder(
+            input_ids= input_ids,
+            attention_mask= attention_mask,
+            token_type_ids= token_type_ids,
+            return_dict= True
+        )
+        encoded_image= self.image_encoder(
+            pixel_values= pixel_values,
+            return_dict= True
+        )
+        combined_output= self.combine(
+            torch.cat(
+                [
+                    encoded_text['pooler_output'],
+                    encoded_image['pooler_output']
+                ],
+                dim= 1
+            )
+        )
+        logits= self.classifier(combined_output)
+        output= {'logits': logits}
+        if labels is not None:
+            loss= self.criterion(logits, labels)
+            output['loss']= loss
+
+        return output