Update app.py

app.py

@@ -87,17 +87,18 @@ def sequence_to_kmer_vector(sequence: str, k: int = 4) -> np.ndarray:
 def calculate_shap_values(model, x_tensor):
     model.eval()
     device = next(model.parameters()).device
-    
+
     try:
-        # Create background as a torch tensor
+        # Create background as a torch tensor (using zeros may be acceptable for DeepExplainer)
         background = torch.zeros((300, x_tensor.shape[1]), device=device)
         explainer = shap.DeepExplainer(model, background)
         shap_values_all = explainer.shap_values(x_tensor)
-        # For binary classification, use the second output and then the first sample
+        # For binary classification, get SHAP for class 1 and first sample
         shap_values = shap_values_all[1][0]
     except Exception as e:
         print(f"DeepExplainer failed, falling back to KernelExplainer: {str(e)}")
-        
+
+        # Define a wrapper that ensures proper input shape and conversion to tensor
         def model_predict(x):
             if not isinstance(x, np.ndarray):
                 x = np.array(x)
@@ -106,24 +107,26 @@ def calculate_shap_values(model, x_tensor):
             with torch.no_grad():
                 tensor_x = torch.tensor(x, dtype=torch.float, device=device)
                 output = model(tensor_x)
-                probs = torch.softmax(output, dim=1)[:, 1]
+                probs = torch.softmax(output, dim=1)[:, 1]  # human probability
                 return probs.cpu().numpy()
-        
-        # Use a numpy background for KernelExplainer
-        background = np.zeros((300, x_tensor.shape[1]))
-        explainer = shap.KernelExplainer(model_predict, background)
+
+        # Instead of using zeros as background, use the input sample repeated 300 times.
         x_numpy = x_tensor.cpu().numpy()
+        background = np.repeat(x_numpy, 300, axis=0)
+
+        explainer = shap.KernelExplainer(model_predict, background)
+        # Increase nsamples for a more robust estimate.
         shap_values = explainer.shap_values(x_numpy, nsamples=1000)
-        # If KernelExplainer returns a list, take its first element.
+        # If a list is returned, select the first element.
         if isinstance(shap_values, list):
             shap_values = shap_values[0]
-    
-    # Get human probability from model prediction
+
+    # Get the human probability from the model output.
     with torch.no_grad():
         output = model(x_tensor)
         probs = torch.softmax(output, dim=1)
         prob_human = probs[0, 1].item()
-    
+
     return np.array(shap_values), prob_human