Upload analysis.py

analysis.py

@@ -36,9 +36,6 @@ def plot_crossattention_weights(target_mask, drug_mask, target_tokenized, drug_t
     if hasattr(drug_mask, 'cpu'):
         drug_mask = drug_mask.cpu()
     
-    # Convert boolean masks
-    target_mask = target_mask.astype(bool)
-    drug_mask = drug_mask.astype(bool)
     logger.info(f"Target mask shape: {target_mask.shape}, Drug mask shape: {drug_mask.shape}")
 
     # Get tokens for unmasked positions
@@ -56,32 +53,33 @@ def plot_crossattention_weights(target_mask, drug_mask, target_tokenized, drug_t
         crossattention_weights = crossattention_weights.cpu()
     
     subset = crossattention_weights[target_mask][:, drug_mask]
-    logger.info(f"Subset shape: {subset.shape}")
-    height, width = subset.shape
+    subset_np = subset.detach().numpy()  # Convert to numpy for matplotlib
     
+    logger.info(f"Subset shape: {subset_np.shape}")
+    height, width = subset_np.shape
+   
     fig, ax = plt.subplots(
         figsize=(width * 0.2 + 2, height * 0.2 + 3),
         dpi=300
     )
-    im = ax.imshow(subset, cmap='hot', interpolation='nearest')
-    
+    im = ax.imshow(subset_np, cmap='hot', interpolation='nearest')
+   
     plt.colorbar(im, ax=ax, orientation='vertical', fraction=0.05, shrink=0.8)
-
     plt.title("Cross-Attention Weights")
     plt.xlabel("Drug Tokens")
     plt.ylabel("Target Tokens")
-    
+   
     # Create vertical labels for drug tokens
     vertical_labels = ['\n'.join(label) for label in drug_token_str]
     plt.xticks(ticks=np.arange(width), labels=vertical_labels)
     plt.yticks(ticks=np.arange(height), labels=target_token_str)
-    
+   
     # Add text annotations
-    max_val = subset.max()
+    max_val = subset_np.max()
     logger.info(f"Max crossattention weight: {max_val}")
     for i in range(height):
         for j in range(width):
-            val = subset[i, j]
+            val = subset_np[i, j]
             if val > max_val / 2:
                 # Extract just the digits after the decimal (no leading '0.')
                 text = f"{val % 1:.2f}"[2:]
@@ -89,14 +87,14 @@ def plot_crossattention_weights(target_mask, drug_mask, target_tokenized, drug_t
                         ha='center', va='center',
                         color="black",
                         fontsize=6)
-    
+   
     # Convert to PIL Image
     buf = io.BytesIO()
     plt.savefig(buf, format='png', bbox_inches='tight', dpi=300)
     buf.seek(0)
     img = Image.open(buf)
     plt.close()
-    
+   
     logger.info("Finished plot_crossattention_weights successfully")
     return img