Update app.py

app.py

@@ -118,12 +118,38 @@ def predict(file_obj):
         kmer_vector = scaler.transform(raw_freq_vector.reshape(1, -1))
         X_tensor = torch.FloatTensor(kmer_vector).to(device)
         
-        # Get feature importance and human probability
-        importance, human_prob = model.get_feature_importance(X_tensor)
-        kmer_importance = importance[0].cpu().numpy()
+        # Calculate final probabilities first
+        with torch.no_grad():
+            output = model(X_tensor)
+            probs = torch.softmax(output, dim=1)
+            human_prob = float(probs[0][1])
+        
+        # Get feature importance using integrated gradients
+        baseline = torch.zeros_like(X_tensor)  # baseline of zeros
+        steps = 50
+        
+        all_importance = []
+        for i in range(steps + 1):
+            alpha = i / steps
+            interpolated = baseline + alpha * (X_tensor - baseline)
+            interpolated.requires_grad_(True)
+            
+            output = model(interpolated)
+            probs = torch.softmax(output, dim=1)
+            human_class = probs[..., 1]
+            
+            if interpolated.grad is not None:
+                interpolated.grad.zero_()
+            human_class.backward()
+            all_importance.append(interpolated.grad.cpu().numpy())
         
-        # Scale importance values relative to the prediction
-        kmer_importance = kmer_importance * human_prob
+        # Average the gradients
+        kmer_importance = np.mean(all_importance, axis=0)[0]
+        # Scale to match probability difference
+        target_diff = human_prob - 0.5  # difference from neutral prediction
+        current_sum = np.sum(kmer_importance)
+        if current_sum != 0:  # avoid division by zero
+            kmer_importance = kmer_importance * (target_diff / current_sum)
         
         # Get top k-mers by absolute importance
         top_k = 10
@@ -157,10 +183,9 @@ def predict(file_obj):
             human_prob = float(probs[0][1])
 
         # Create SHAP explanation
-        # We'll use the actual probabilities for alignment
         explanation = shap.Explanation(
             values=np.array(top_values),
-            base_values=human_prob,  # Use actual prediction as base
+            base_values=0.5,  # Start from neutral prediction
             data=np.array([
                 raw_freq_vector[kmer_dict[feat]] if feat != "Others" 
                 else np.sum(raw_freq_vector[others_mask]) 
@@ -168,7 +193,7 @@ def predict(file_obj):
             ]),
             feature_names=top_features
         )
-        explanation.expected_value = human_prob  # Match the actual prediction
+        explanation.expected_value = 0.5  # Start from neutral prediction
 
         # Create waterfall plot
         plt.figure(figsize=(10, 6))