Update app.py

app.py

@@ -88,59 +88,59 @@ def sequence_to_kmer_vector(sequence: str, k: int = 4) -> np.ndarray:
 def calculate_shap_values(model, x_tensor, baseline=None, steps=50):
     """
     Calculate feature attributions using Integrated Gradients.
-
+    
     Args:
         model: A PyTorch model.
         x_tensor: Input tensor of shape (1, num_features).
         baseline: Tensor of the same shape as x_tensor to use as the reference.
                   If None, defaults to a tensor of zeros.
         steps: Number of steps in the Riemann approximation of the integral.
-
+    
     Returns:
         attributions: A numpy array of shape (num_features,) with feature attributions.
-        baseline_prob: The model's predicted probability for the target class (human)
-                       when using the baseline input.
+        full_prob: The model's predicted probability for the target class (human)
+                   when using the actual input.
     """
     model.eval()
     if baseline is None:
         baseline = torch.zeros_like(x_tensor)
-
+    
+    # Compute the model's prediction for the full input.
+    with torch.no_grad():
+        full_output = model(x_tensor)
+        full_probs = torch.softmax(full_output, dim=1)
+        full_prob = full_probs[0, 1].item()  # Probability for 'human'
+    
     # Generate interpolated inputs between the baseline and the actual input.
     scaled_inputs = [
         baseline + (float(i) / steps) * (x_tensor - baseline)
         for i in range(steps + 1)
     ]
-    scaled_inputs = torch.cat(scaled_inputs, dim=0)  # shape: (steps+1, num_features)
+    scaled_inputs = torch.cat(scaled_inputs, dim=0)  # Shape: (steps+1, num_features)
     scaled_inputs.requires_grad = True
 
     # Forward pass: compute model outputs for all interpolated inputs.
-    outputs = model(scaled_inputs)  # shape: (steps+1, num_classes)
-    probs = torch.softmax(outputs, dim=1)[:, 1]  # probability for the 'human' class
+    outputs = model(scaled_inputs)  # Shape: (steps+1, num_classes)
+    probs = torch.softmax(outputs, dim=1)[:, 1]  # Probability for 'human'
 
-    # Backward pass: compute gradients of the probability with respect to the inputs.
+    # Backward pass: compute gradients of the probability with respect to inputs.
     grads = torch.autograd.grad(
         outputs=probs,
         inputs=scaled_inputs,
         grad_outputs=torch.ones_like(probs),
         create_graph=False,
         retain_graph=False
-    )[0]  # shape: (steps+1, num_features)
+    )[0]  # Shape: (steps+1, num_features)
 
     # Approximate the integral using the trapezoidal rule.
-    # Compute the average gradient between consecutive steps.
-    avg_grads = (grads[:-1] + grads[1:]) / 2.0
-    # Average the gradients over all steps.
-    integrated_grad = avg_grads.mean(dim=0, keepdim=True)  # shape: (1, num_features)
+    avg_grads = (grads[:-1] + grads[1:]) / 2.0  # Average gradient between steps.
+    integrated_grad = avg_grads.mean(dim=0, keepdim=True)  # Mean over all steps.
 
-    # Scale the integrated gradients by the difference between the input and the baseline.
-    attributions = (x_tensor - baseline) * integrated_grad  # shape: (1, num_features)
+    # Scale the integrated gradients by the difference between the input and baseline.
+    attributions = (x_tensor - baseline) * integrated_grad
 
-    # Compute the baseline probability (for reference)
-    with torch.no_grad():
-        baseline_output = model(baseline)
-        baseline_prob = torch.softmax(baseline_output, dim=1)[0, 1].item()
+    return attributions.squeeze().cpu().numpy(), full_prob
 
-    return attributions.squeeze().cpu().numpy(), baseline_prob
 
 
 ###############################################################################