Update app.py

app.py

@@ -150,12 +150,17 @@ def predict(file_obj):
         top_features.append("Others")
         top_values.append(others_sum)
         
+        # Calculate final probabilities first
+        with torch.no_grad():
+            output = model(X_tensor)
+            probs = torch.softmax(output, dim=1)
+            human_prob = float(probs[0][1])
+
         # Create SHAP explanation
-        # Set base_value to 0.5 (neutral prediction)
-        # Values represent the push towards human (>0.5) or non-human (<0.5)
+        # We'll use the actual probabilities for alignment
         explanation = shap.Explanation(
             values=np.array(top_values),
-            base_values=0.5,  # Start from neutral prediction
+            base_values=human_prob,  # Use actual prediction as base
             data=np.array([
                 raw_freq_vector[kmer_dict[feat]] if feat != "Others" 
                 else np.sum(raw_freq_vector[others_mask]) 
@@ -163,7 +168,7 @@ def predict(file_obj):
             ]),
             feature_names=top_features
         )
-        explanation.expected_value = 0.5
+        explanation.expected_value = human_prob  # Match the actual prediction
 
         # Create waterfall plot
         plt.figure(figsize=(10, 6))
@@ -172,7 +177,7 @@ def predict(file_obj):
             show=False,
             max_display=11  # Show all features including "Others"
         )
-        plt.title(f"Impact on prediction (>0.5 pushes toward human, <0.5 toward non-human)")
+        plt.title(f"Feature contributions to human probability (final prob: {human_prob:.3f})")
         
         # Save plot
         buf = io.BytesIO()
@@ -220,4 +225,4 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch(share=True)
+    iface.launch(share=True)