[FIX]UI

app.py

@@ -74,16 +74,27 @@ class LateFusionModel(nn.Module):
         super(LateFusionModel, self).__init__()
         self.image_model = image_model
         self.text_model = text_model
-        self.image_weight = nn.Parameter(torch.tensor(0.5))
-        self.text_weight = nn.Parameter(torch.tensor(0.5))
+        # MLP fusion layer (matching saved model structure)
+        # Structure: Linear(2, hidden) -> ReLU -> Dropout -> Linear(hidden, 1)
+        hidden_dim = 64  # Adjust if needed based on saved model
+        self.fusion_mlp = nn.Sequential(
+            nn.Linear(2, hidden_dim),  # layer 0
+            nn.ReLU(),                 # layer 1 (no params)
+            nn.Dropout(0.1),           # layer 2 (no params)
+            nn.Linear(hidden_dim, 1)   # layer 3
+        )
 
     def forward(self, images, input_ids, attention_mask):
         with torch.no_grad():
             image_logits = self.image_model(images).squeeze(1)
             text_logits = self.text_model(input_ids=input_ids, attention_mask=attention_mask).logits.squeeze(1)
 
-        weights = torch.softmax(torch.stack([self.image_weight, self.text_weight]), dim=0)
-        fused_logits = weights[0] * image_logits + weights[1] * text_logits
+        # Stack logits and pass through MLP
+        stacked_logits = torch.stack([image_logits, text_logits], dim=1)
+        fused_logits = self.fusion_mlp(stacked_logits).squeeze(1)
+        
+        # For compatibility, create dummy weights
+        weights = torch.tensor([0.5, 0.5], device=fused_logits.device)
 
         return fused_logits, image_logits, text_logits, weights
 
@@ -91,7 +102,11 @@ class LateFusionModel(nn.Module):
 # Create model architecture first
 image_model_for_fusion = models.efficientnet_b3(weights=models.EfficientNet_B3_Weights.DEFAULT)
 num_features = image_model_for_fusion.classifier[1].in_features
-image_model_for_fusion.classifier = nn.Linear(num_features, 1)
+# Match saved model structure: classifier.1 instead of classifier
+image_model_for_fusion.classifier = nn.Sequential(
+    nn.Dropout(p=0.3, inplace=True),
+    nn.Linear(num_features, 1)
+)
 
 text_base_model = AutoModel.from_pretrained('indobenchmark/indobert-base-p1')
 text_model = TextModelWithClassifier(text_base_model)
@@ -404,14 +419,23 @@ def predict_single_url(url):
         }
         
         confidence = gambling_prob if is_gambling else non_gambling_prob
-        image_weight = weights[0].item()
-        text_weight = weights[1].item()
+        
+        # Calculate relative contribution (approximation for MLP fusion)
+        image_contrib = abs(image_probs[0].item() - 0.5)
+        text_contrib = abs(text_probs[0].item() - 0.5)
+        total_contrib = image_contrib + text_contrib
+        if total_contrib > 0:
+            image_weight = image_contrib / total_contrib
+            text_weight = text_contrib / total_contrib
+        else:
+            image_weight = 0.5
+            text_weight = 0.5
         
         confidence_md = f"**Confidence:** {confidence:.1%}\n\n**Model Used:** Fusion Model (Image + Text)\n\n**Prediction:** {'🟥 Gambling' if is_gambling else '🟩 Non-Gambling'}"
         
-        model_info = f"""**Model Type:** Fusion Model
-**Image Model:** EfficientNet-B3 (Weight: {image_weight:.1%})
-**Text Model:** IndoBERT (Weight: {text_weight:.1%})
+        model_info = f"""**Model Type:** Fusion Model (MLP)
+**Image Model:** EfficientNet-B3
+**Text Model:** IndoBERT
 
 **Individual Predictions:**
 - 🖼️ Image Model: {image_probs[0].item():.1%}
@@ -610,9 +634,9 @@ with gr.Blocks(theme=gr.themes.Soft(), css=custom_css, title="Gambling Website D
                 with gr.Column():
                     file_input = gr.File(
                         label="Upload URL File (.txt)",
-                        file_types=[".txt"],
-                        info="Upload a .txt file with one URL per line"
+                        file_types=[".txt"]
                     )
+                    gr.Markdown("💡 **Tip:** Upload a .txt file with one URL per line")
                     batch_predict_button = gr.Button(
                         "🚀 Process Batch",
                         variant="primary",