fixed pytorch

model_inference.py

@@ -19,26 +19,48 @@ def ensure_model_file(filename: str):
     return path
 
 
-def load_model(path):
+import os
+import torch
+from model_architecture import EnsembleBertBiLSTMRegressor
+
+def load_model(path: str = "assets/best_checkpoint_regression.pt"):
+    """
+    Load the trained EnsembleBertBiLSTMRegressor model using saved checkpoint weights.
+    Supports CPU/GPU execution.
+    """
+
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"❌ Model checkpoint not found at: {path}")
+
+    print(f"✅ Loading model weights from: {path}")
     checkpoint = torch.load(path, map_location="cpu")
-    
-    # Recreate the same model architecture
+
+    # ✅ Recreate model architecture (must match training exactly!)
     model = EnsembleBertBiLSTMRegressor(
         model_name_mcq="microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract",
         model_name_clinical="emilyalsentzer/Bio_ClinicalBERT",
         hidden_dim=768,
-        extra_dim=67   # e.g. 10 if you have 10 engineered + categorical features
+        extra_dim=67   # Adjust if your engineered features size differs
     )
 
-    # Load saved weights
-    model.load_state_dict(checkpoint["model_state"])
-    model.eval()
+    # ✅ Load weights into model
+    if "model_state" in checkpoint:
+        model.load_state_dict(checkpoint["model_state"])
+    elif "state_dict" in checkpoint:   # support alternative saving formats
+        model.load_state_dict(checkpoint["state_dict"])
+    else:
+        raise KeyError("❌ No 'model_state' or 'state_dict' found in checkpoint")
 
+    # ✅ Set eval mode and move to device
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
+    model.eval()
+
+    print(f"✅ Model loaded successfully on device: {device}")
     return model, device
 
 
+
 def predict_from_input(data, model, device, tok_mcq, tok_clin, encoder, scaler):
     """
     Predict difficulty and discrimination index for a single MCQ item.