Create model.py

model.py

@@ -0,0 +1,56 @@
+# model.py
+
+import torch
+import torch.nn as nn
+from transformers import AutoModelForSeq2SeqLM
+from torchvision import models
+
+class ImageToTextProjector(nn.Module):
+    def __init__(self, image_embedding_dim, text_embedding_dim):
+        super(ImageToTextProjector, self).__init__()
+        self.fc = nn.Linear(image_embedding_dim, text_embedding_dim)
+        self.activation = nn.ReLU()
+        self.dropout = nn.Dropout(p=0.5)
+
+    def forward(self, x):
+        x = self.fc(x)
+        x = self.activation(x)
+        x = self.dropout(x)
+        return x
+
+class CombinedModel(nn.Module):
+    def __init__(self, video_model, report_generator, num_classes, projector):
+        super(CombinedModel, self).__init__()
+        self.video_model = video_model
+        self.report_generator = report_generator
+        self.classifier = nn.Linear(512, num_classes)
+        self.projector = projector
+        self.dropout = nn.Dropout(p=0.5)
+
+    def forward(self, images, labels=None):
+        video_embeddings = self.video_model(images)
+        video_embeddings = self.dropout(video_embeddings)
+        class_outputs = self.classifier(video_embeddings)
+        projected_embeddings = self.projector(video_embeddings)
+        encoder_inputs = projected_embeddings.unsqueeze(1)
+
+        if labels is not None:
+            outputs = self.report_generator(
+                inputs_embeds=encoder_inputs,
+                labels=labels
+            )
+            gen_loss = outputs.loss
+            generated_report = None
+        else:
+            generated_report_ids = self.report_generator.generate(
+                inputs_embeds=encoder_inputs,
+                max_length=512,
+                num_beams=4,
+                early_stopping=True
+            )
+            generated_report = report_generator_tokenizer.batch_decode(
+                generated_report_ids, skip_special_tokens=True
+            )
+            gen_loss = None
+
+        return class_outputs, generated_report, gen_loss