initial commit

app.py

@@ -12,16 +12,16 @@ class MultimodalRiskBehaviorModel(nn.Module):
     def __init__(self, text_model_name="bert-base-uncased", hidden_dim=512, dropout=0.3):
         super(MultimodalRiskBehaviorModel, self).__init__()
 
-        # Use AutoModelForSequenceClassification for classification tasks
+        # Text model using AutoModelForSequenceClassification
         self.text_model_name = text_model_name
         self.text_model = AutoModelForSequenceClassification.from_pretrained(text_model_name, num_labels=1)
         
-        # Visual model initialization with ResNet50
+        # Visual model (ResNet50)
         self.visual_model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
         visual_feature_dim = self.visual_model.fc.in_features
         self.visual_model.fc = nn.Identity()
 
-        # Fusion and classification layers
+        # Fusion and classification layer setup
         text_feature_dim = self.text_model.config.hidden_size
         self.fc1 = nn.Linear(text_feature_dim + visual_feature_dim, hidden_dim)
         self.dropout = nn.Dropout(dropout)
@@ -31,22 +31,32 @@ class MultimodalRiskBehaviorModel(nn.Module):
         input_ids = encoding['input_ids'].squeeze(1).to(device)
         attention_mask = encoding['attention_mask'].squeeze(1).to(device)
 
-        # Text and visual features extraction
-        text_features = self.text_model(input_ids=input_ids, attention_mask=attention_mask).logits.squeeze(-1)
+        # Extract text and visual features
+        text_features = self.text_model(input_ids=input_ids, attention_mask=attention_mask).logits
         frames = frames.to(device)
         
         batch_size, num_frames, channels, height, width = frames.size()
         frames = frames.view(batch_size * num_frames, channels, height, width)
         visual_features = self.visual_model(frames)
         visual_features = visual_features.view(batch_size, num_frames, -1).mean(dim=1)
-        
-        # Combine features and classify
+
+        # Combine and classify
         combined_features = torch.cat((text_features, visual_features), dim=1)
         x = self.dropout(torch.relu(self.fc1(combined_features)))
         output = torch.sigmoid(self.fc2(x))
 
         return output
 
+    def save_pretrained(self, save_directory):
+        os.makedirs(save_directory, exist_ok=True)
+        torch.save(self.state_dict(), os.path.join(save_directory, 'pytorch_model.bin'))
+        config = {
+            "text_model_name": self.text_model_name,
+            "hidden_dim": self.fc1.out_features
+        }
+        with open(os.path.join(save_directory, 'config.json'), 'w') as f:
+            json.dump(config, f)
+
     @classmethod
     def from_pretrained(cls, load_directory, map_location=None):
         if os.path.exists(load_directory):