Update app.py

app.py

@@ -4,7 +4,92 @@ from PIL import Image
 from torchvision import transforms
 from transformers import T5Tokenizer, ViTFeatureExtractor
 
+class Encoder(nn.Module):
+    def __init__(self, pretrained_model):
+        """
+        Implements the Encoder."
+
+        Args:
+            pretrained_model (str): name of the pretrained model
+
+        """
+
+        super(Encoder, self).__init__()
+
+        self.encoder = ViTModel.from_pretrained(pretrained_model)
+
+    def forward(self, input):
+        out  = self.encoder(pixel_values = input)
+
+        return out
+
+class Decoder(nn.Module):
+    def __init__(self, pretrained_model, encoder_modeldim):
+        """
+        Implements the Decoder."
+
+        Args:
+            pretrained_model (str): name of the pretrained model
+
+        """
+
+        super(Decoder, self).__init__()
+
+        self.decoder = T5ForConditionalGeneration.from_pretrained(pretrained_model)
+        self.linear = nn.Linear(self.decoder.model_dim, encoder_modeldim, bias = False)
+        self.encoder_modeldim = encoder_modeldim
+
+    def forward(self, output_encoder, targets, decoder_ids=None):
+
+        if self.decoder.model_dim!=self.encoder_modeldim:
+            print(f"Changed model hidden dimension from {self.encoder_modeldim} to {self.decoder.model_dim}")
+            output_encoder = self.linear(output_encoder)
+            print(output_encoder.shape)
+
+        # Validation/Testing
+        if decoder_ids is not None:
+            out = self.decoder(encoder_outputs=output_encoder, decoder_input_ids=decoder_ids)
+
+        # Training
+        else:
+            out = self.decoder(encoder_outputs=output_encoder, labels=targets)
+
+        return out
+
+class EncoderDecoder(nn.Module):
+    def __init__(self, pretrained_model: Tuple[str], encoder_dmodel=768, eos_token_id=None, pad_token_id=None):
+        """
+        Implements a model that combines MyEncoder and MyDecoder."
+
+        Args:
+            pretrained_model (tuple): name of the pretrained model
+            encoder_dmodel (int):  hidden dimension of the encoder output
+            eos_token_id (torch.long): token used for end of sentence
+            pad_token_id (torch.long): token used for padding
+
+        """
+
+        super(EncoderDecoder, self).__init__()
+        self.eos_token_id = eos_token_id
+        self.pad_token_id = pad_token_id
+        self.encoder = Encoder(pretrained_model[0])
+        self.encoder_dmodel = encoder_dmodel
+
+        # Freeze parameters from encoder
+        #for p in self.encoder.parameters():
+        #    p.requires_grad=False
+
+        self.decoder = Decoder(pretrained_model[1], self.encoder_dmodel)
+        self.decoder_start_token_id = self.decoder.decoder.config.decoder_start_token_id
+
+    def forward(self, images = None, targets = None, decoder_ids = None):
+        output_encoder  = self.encoder(images)
+        out = self.decoder(output_encoder, targets, decoder_ids)
+
+        return out
+
 # Model loading and setting up the device
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = torch.load("model_vit_ai.pt", map_location=device)
 model.to(device)