Upload analysis.py

analysis.py

@@ -32,11 +32,9 @@ def plot_crossattention_weights(target_mask, drug_mask, target_tokenized, drug_t
 
     # Convert masks to numpy arrays if they're tensors
     if hasattr(target_mask, 'cpu'):
-        logger.info("Converting target_mask tensor to numpy")
-        target_mask = target_mask.cpu().numpy()
+        target_mask = target_mask.cpu()
     if hasattr(drug_mask, 'cpu'):
-        logger.info("Converting drug_mask tensor to numpy")
-        drug_mask = drug_mask.cpu().numpy()
+        drug_mask = drug_mask.cpu()
     
     # Convert boolean masks
     target_mask = target_mask.astype(bool)
@@ -45,23 +43,17 @@ def plot_crossattention_weights(target_mask, drug_mask, target_tokenized, drug_t
 
     # Get tokens for unmasked positions
     tokens_input = target_tokenized["input_ids"][0][target_mask]
-    if hasattr(tokens_input, 'cpu'):
-        logger.info("Converting target tokens tensor to numpy")
-        tokens_input = tokens_input.cpu().numpy()
+
     target_token_str = target_tokenizer.convert_ids_to_tokens(tokens_input)
-    logger.info(f"Target tokens: {target_token_str}")
 
     tokens_input = drug_tokenized["input_ids"][0][drug_mask]
-    if hasattr(tokens_input, 'cpu'):
-        logger.info("Converting drug tokens tensor to numpy")
-        tokens_input = tokens_input.cpu().numpy()
     drug_token_str = drug_tokenizer.convert_ids_to_tokens(tokens_input)
+
     logger.info(f"Drug tokens: {drug_token_str}")
 
     # Extract subset of attention weights
     if hasattr(crossattention_weights, 'cpu'):
-        logger.info("Converting crossattention_weights tensor to numpy")
-        crossattention_weights = crossattention_weights.cpu().numpy()
+        crossattention_weights = crossattention_weights.cpu()
     
     subset = crossattention_weights[target_mask][:, drug_mask]
     logger.info(f"Subset shape: {subset.shape}")