Upload app.py

app.py

@@ -158,182 +158,182 @@ class DrugTargetInteractionApp:
             logger.error(f"Prediction error: {str(e)}")
             return f"Error during prediction: {str(e)}"
 
-def visualize_interaction(self, target_sequence, drug_smiles):
-    """
-    Generate visualization images for drug-target interaction
-    
-    Args:
-        target_sequence (str): RNA sequence
-        drug_smiles (str): Drug SMILES notation
+    def visualize_interaction(self, target_sequence, drug_smiles):
+        """
+        Generate visualization images for drug-target interaction
         
-    Returns:
-        tuple: (cross_attention_image, raw_contribution_image, normalized_contribution_image, status_message)
-    """
-    if self.model is None:
-        return None, None, None, "Error: Model not loaded. Please load a model first."
-    
-    try:
-        # Tokenize inputs
-        target_inputs = self.target_tokenizer(
-            target_sequence,
-            padding="max_length", 
-            truncation=True, 
-            max_length=512,
-            return_tensors="pt"
-        ).to(self.device)
-        
-        drug_inputs = self.drug_tokenizer(
-            drug_smiles,
-            padding="max_length", 
-            truncation=True, 
-            max_length=512,
-            return_tensors="pt"
-        ).to(self.device)
-        
-        # Enable interpretation mode
-        self.model.INTERPR_ENABLE_MODE()
-
-        # Make prediction and extract visualization data
-        with torch.no_grad():
-            prediction = self.model(target_inputs, drug_inputs)
-            
-            # Unscale if scaler exists
-            if self.model.scaler is not None:
-                prediction = self.model.unscale(prediction)
+        Args:
+            target_sequence (str): RNA sequence
+            drug_smiles (str): Drug SMILES notation
             
-            prediction_value = prediction.cpu().numpy()[0][0]
-            
-            # Extract data needed for visualizations
-            presum_values = self.model.model.presum_layer  # Shape: (1, seq_len)
-            cross_attention_weights = self.model.model.crossattention_weights  # Shape: (batch, heads, seq_len, seq_len)
-            
-            # Get model parameters for scaling
-            w = self.model.model.w.squeeze(1)
-            b = self.model.model.b
-            scaler = self.model.model.scaler
-            
-        logger.info(f"Target inputs shape: {target_inputs['input_ids'].shape}")
-        logger.info(f"Drug inputs shape: {drug_inputs['input_ids'].shape}")
+        Returns:
+            tuple: (cross_attention_image, raw_contribution_image, normalized_contribution_image, status_message)
+        """
+        if self.model is None:
+            return None, None, None, "Error: Model not loaded. Please load a model first."
         
-        # Generate visualizations
         try:
-            # 1. Cross-attention heatmap
-            cross_attention_img = None
-            logger.info(f"Cross-attention weights type: {type(cross_attention_weights)}")
-            if cross_attention_weights is not None:
-                logger.info(f"Cross-attention weights shape: {cross_attention_weights.shape if hasattr(cross_attention_weights, 'shape') else 'No shape attr'}")
-                
-                try:
-                    cross_attn_matrix = cross_attention_weights[0, 0]
-                        
-                    if cross_attn_matrix is not None:
-                        logger.info(f"Extracted cross-attention matrix shape: {cross_attn_matrix.shape}")
-                        logger.info(f"Target attention mask shape: {target_inputs['attention_mask'].shape}")
-                        logger.info(f"Drug attention mask shape: {drug_inputs['attention_mask'].shape}")
-
+            # Tokenize inputs
+            target_inputs = self.target_tokenizer(
+                target_sequence,
+                padding="max_length", 
+                truncation=True, 
+                max_length=512,
+                return_tensors="pt"
+            ).to(self.device)
+            
+            drug_inputs = self.drug_tokenizer(
+                drug_smiles,
+                padding="max_length", 
+                truncation=True, 
+                max_length=512,
+                return_tensors="pt"
+            ).to(self.device)
+            
+            # Enable interpretation mode
+            self.model.INTERPR_ENABLE_MODE()
 
-                        cross_attention_img = plot_crossattention_weights(
-                            target_inputs["attention_mask"][0],
-                            drug_inputs["attention_mask"][0],
-                            target_inputs,
-                            drug_inputs,
-                            cross_attn_matrix,
-                            self.target_tokenizer,
-                            self.drug_tokenizer
-                        )
-                    else:
-                        logger.warning("Could not extract valid cross-attention matrix")
-                        
-                except (IndexError, TypeError, AttributeError) as e:
-                    logger.warning(f"Error extracting cross-attention matrix: {str(e)}")
-                    cross_attn_matrix = None
-            else:
-                logger.warning("Cross-attention weights are None")
+            # Make prediction and extract visualization data
+            with torch.no_grad():
+                prediction = self.model(target_inputs, drug_inputs)
                 
-        except Exception as e:
-            logger.error(f"Cross-attention visualization error: {str(e)}")
-            cross_attention_img = None
-        
-        try:
-            # 2. Normalized contribution visualization (always generate)
-            normalized_img = None
-            if presum_values is not None:
-                normalized_img = plot_presum(
-                    target_inputs,
-                    presum_values.detach(),  # Detach the tensor
-                    scaler,
-                    w.detach(),  # Detach the tensor
-                    b.detach(),  # Detach the tensor
-                    self.target_tokenizer,
-                    raw_affinities=False
-                )
-            else:
-                logger.warning("Presum values are None")
+                # Unscale if scaler exists
+                if self.model.scaler is not None:
+                    prediction = self.model.unscale(prediction)
                 
-        except Exception as e:
-            logger.error(f"Normalized contribution visualization error: {str(e)}")
-            normalized_img = None
-        
-        try:
-            # 3. Raw contribution visualization (only if pKd > 0)
-            raw_img = None
-            if prediction_value > 0 and presum_values is not None:
-                raw_img = plot_presum(
-                    target_inputs,
-                    presum_values.detach(),  # Detach the tensor
-                    scaler,
-                    w.detach(),  # Detach the tensor
-                    b.detach(),  # Detach the tensor
-                    self.target_tokenizer,
-                    raw_affinities=True
-                )
-            else:
-                if prediction_value <= 0:
-                    logger.info("Skipping raw affinities visualization as pKd <= 0")
-                if presum_values is None:
-                    logger.warning("Cannot generate raw visualization: presum values are None")
+                prediction_value = prediction.cpu().numpy()[0][0]
+                
+                # Extract data needed for visualizations
+                presum_values = self.model.model.presum_layer  # Shape: (1, seq_len)
+                cross_attention_weights = self.model.model.crossattention_weights  # Shape: (batch, heads, seq_len, seq_len)
+                
+                # Get model parameters for scaling
+                w = self.model.model.w.squeeze(1)
+                b = self.model.model.b
+                scaler = self.model.model.scaler
+                
+            logger.info(f"Target inputs shape: {target_inputs['input_ids'].shape}")
+            logger.info(f"Drug inputs shape: {drug_inputs['input_ids'].shape}")
+            
+            # Generate visualizations
+            try:
+                # 1. Cross-attention heatmap
+                cross_attention_img = None
+                logger.info(f"Cross-attention weights type: {type(cross_attention_weights)}")
+                if cross_attention_weights is not None:
+                    logger.info(f"Cross-attention weights shape: {cross_attention_weights.shape if hasattr(cross_attention_weights, 'shape') else 'No shape attr'}")
                     
-        except Exception as e:
-            logger.error(f"Raw contribution visualization error: {str(e)}")
-            raw_img = None
-        
-        # Disable interpretation mode after use
-        self.model.INTERPR_DISABLE_MODE()
-        
-        # Create placeholder images if generation failed
-        if cross_attention_img is None:
-            cross_attention_img = create_placeholder_image(
-                text="Cross-Attention Heatmap\nFailed to generate"
-            )
-        if normalized_img is None:
-            normalized_img = create_placeholder_image(
-                text="Normalized Contribution\nFailed to generate"
-            )
-        if raw_img is None and prediction_value > 0:
-            raw_img = create_placeholder_image(
-                text="Raw Contribution\nFailed to generate"
-            )
-        elif raw_img is None:
-            raw_img = create_placeholder_image(
-                text="Raw Contribution\nSkipped (pKd ≤ 0)"
-            )
+                    try:
+                        cross_attn_matrix = cross_attention_weights[0, 0]
+                            
+                        if cross_attn_matrix is not None:
+                            logger.info(f"Extracted cross-attention matrix shape: {cross_attn_matrix.shape}")
+                            logger.info(f"Target attention mask shape: {target_inputs['attention_mask'].shape}")
+                            logger.info(f"Drug attention mask shape: {drug_inputs['attention_mask'].shape}")
+
+
+                            cross_attention_img = plot_crossattention_weights(
+                                target_inputs["attention_mask"][0],
+                                drug_inputs["attention_mask"][0],
+                                target_inputs,
+                                drug_inputs,
+                                cross_attn_matrix,
+                                self.target_tokenizer,
+                                self.drug_tokenizer
+                            )
+                        else:
+                            logger.warning("Could not extract valid cross-attention matrix")
+                            
+                    except (IndexError, TypeError, AttributeError) as e:
+                        logger.warning(f"Error extracting cross-attention matrix: {str(e)}")
+                        cross_attn_matrix = None
+                else:
+                    logger.warning("Cross-attention weights are None")
+                    
+            except Exception as e:
+                logger.error(f"Cross-attention visualization error: {str(e)}")
+                cross_attention_img = None
+            
+            try:
+                # 2. Normalized contribution visualization (always generate)
+                normalized_img = None
+                if presum_values is not None:
+                    normalized_img = plot_presum(
+                        target_inputs,
+                        presum_values.detach(),  # Detach the tensor
+                        scaler,
+                        w.detach(),  # Detach the tensor
+                        b.detach(),  # Detach the tensor
+                        self.target_tokenizer,
+                        raw_affinities=False
+                    )
+                else:
+                    logger.warning("Presum values are None")
+                    
+            except Exception as e:
+                logger.error(f"Normalized contribution visualization error: {str(e)}")
+                normalized_img = None
             
-        status_msg = f"Predicted Binding Affinity: {prediction_value:.4f}"
-        if prediction_value <= 0:
-            status_msg += " (Raw contribution visualization skipped due to non-positive pKd)"
-        if cross_attention_weights is None:
-            status_msg += " (Cross-attention visualization failed: weights not available)"
+            try:
+                # 3. Raw contribution visualization (only if pKd > 0)
+                raw_img = None
+                if prediction_value > 0 and presum_values is not None:
+                    raw_img = plot_presum(
+                        target_inputs,
+                        presum_values.detach(),  # Detach the tensor
+                        scaler,
+                        w.detach(),  # Detach the tensor
+                        b.detach(),  # Detach the tensor
+                        self.target_tokenizer,
+                        raw_affinities=True
+                    )
+                else:
+                    if prediction_value <= 0:
+                        logger.info("Skipping raw affinities visualization as pKd <= 0")
+                    if presum_values is None:
+                        logger.warning("Cannot generate raw visualization: presum values are None")
+                        
+            except Exception as e:
+                logger.error(f"Raw contribution visualization error: {str(e)}")
+                raw_img = None
             
-        return cross_attention_img, raw_img, normalized_img, status_msg
-        
-    except Exception as e:
-        logger.error(f"Visualization error: {str(e)}")
-        # Make sure to disable interpretation mode even if there's an error
-        try:
+            # Disable interpretation mode after use
             self.model.INTERPR_DISABLE_MODE()
-        except:
-            pass
-        return None, None, None, f"Error during visualization: {str(e)}"
+            
+            # Create placeholder images if generation failed
+            if cross_attention_img is None:
+                cross_attention_img = create_placeholder_image(
+                    text="Cross-Attention Heatmap\nFailed to generate"
+                )
+            if normalized_img is None:
+                normalized_img = create_placeholder_image(
+                    text="Normalized Contribution\nFailed to generate"
+                )
+            if raw_img is None and prediction_value > 0:
+                raw_img = create_placeholder_image(
+                    text="Raw Contribution\nFailed to generate"
+                )
+            elif raw_img is None:
+                raw_img = create_placeholder_image(
+                    text="Raw Contribution\nSkipped (pKd ≤ 0)"
+                )
+                
+            status_msg = f"Predicted Binding Affinity: {prediction_value:.4f}"
+            if prediction_value <= 0:
+                status_msg += " (Raw contribution visualization skipped due to non-positive pKd)"
+            if cross_attention_weights is None:
+                status_msg += " (Cross-attention visualization failed: weights not available)"
+                
+            return cross_attention_img, raw_img, normalized_img, status_msg
+            
+        except Exception as e:
+            logger.error(f"Visualization error: {str(e)}")
+            # Make sure to disable interpretation mode even if there's an error
+            try:
+                self.model.INTERPR_DISABLE_MODE()
+            except:
+                pass
+            return None, None, None, f"Error during visualization: {str(e)}"
 
 
 # Initialize the app