[UPDATE]Model

app.py

@@ -60,11 +60,11 @@ print("OCR reader initialized.")
 class TextModelWithClassifier(nn.Module):
     def __init__(self, base_model):
         super(TextModelWithClassifier, self).__init__()
-        self.base_model = base_model
+        self.bert = base_model  # Use 'bert' to match saved state_dict keys
         self.classifier = nn.Linear(base_model.config.hidden_size, 1)
     
     def forward(self, input_ids, attention_mask):
-        outputs = self.base_model(input_ids=input_ids, attention_mask=attention_mask)
+        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
         pooled_output = outputs.pooler_output if hasattr(outputs, 'pooler_output') else outputs.last_hidden_state[:, 0]
         logits = self.classifier(pooled_output)
         return type('Output', (), {'logits': logits})()
@@ -101,13 +101,27 @@ fusion_model = LateFusionModel(image_model_for_fusion, text_model)
 # Load state_dict
 model_path = "models/best_mlp_fusion_model_state_dict.pt"
 if os.path.exists(model_path):
-    fusion_model.load_state_dict(torch.load(model_path, map_location=device))
-    print("Fusion model loaded from local state_dict successfully!")
+    state_dict = torch.load(model_path, map_location=device)
+    try:
+        fusion_model.load_state_dict(state_dict, strict=True)
+        print("Fusion model loaded from local state_dict successfully!")
+    except RuntimeError as e:
+        print(f"Warning: Some keys didn't match. Trying with strict=False...")
+        print(f"Error details: {str(e)[:500]}")
+        fusion_model.load_state_dict(state_dict, strict=False)
+        print("Fusion model loaded with strict=False (some keys may be missing)")
 else:
     print("Fusion model not found locally. Downloading from Hugging Face Hub...")
     model_path = hf_hub_download(repo_id="azzandr/gambling-fusion-model", filename="best_mlp_fusion_model_state_dict.pt")
-    fusion_model.load_state_dict(torch.load(model_path, map_location=device))
-    print("Fusion model downloaded and loaded successfully!")
+    state_dict = torch.load(model_path, map_location=device)
+    try:
+        fusion_model.load_state_dict(state_dict, strict=True)
+        print("Fusion model downloaded and loaded successfully!")
+    except RuntimeError as e:
+        print(f"Warning: Some keys didn't match. Trying with strict=False...")
+        print(f"Error details: {str(e)[:500]}")
+        fusion_model.load_state_dict(state_dict, strict=False)
+        print("Fusion model loaded with strict=False (some keys may be missing)")
 
 fusion_model.to(device)
 fusion_model.eval()