Upload 2 files

analysis.py

@@ -0,0 +1,211 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import os
+from matplotlib.colors import LinearSegmentedColormap
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from PIL import Image
+import io
+
+def plot_crossattention_weights(target_mask, drug_mask, target_tokenized, drug_tokenized, 
+                               crossattention_weights, target_tokenizer, drug_tokenizer):
+    """
+    Plots the cross-attention weights for a given drug-target pair, only considering unmasked tokens.
+    
+    Parameters:
+        target_mask (np.ndarray): Boolean mask for target tokens.
+        drug_mask (np.ndarray): Boolean mask for drug tokens.
+        target_tokenized (dict): Tokenized target sequence.
+        drug_tokenized (dict): Tokenized drug sequence.
+        crossattention_weights (np.ndarray): The cross-attention weights.
+        target_tokenizer: Target tokenizer instance.
+        drug_tokenizer: Drug tokenizer instance.
+        
+    Returns:
+        PIL.Image: The generated attention heatmap image.
+    """
+    # Convert masks to numpy arrays if they're tensors
+    if hasattr(target_mask, 'cpu'):
+        target_mask = target_mask.cpu().numpy()
+    if hasattr(drug_mask, 'cpu'):
+        drug_mask = drug_mask.cpu().numpy()
+    
+    # Convert boolean masks
+    target_mask = target_mask.astype(bool)
+    drug_mask = drug_mask.astype(bool)
+    
+    # Get tokens for unmasked positions
+    tokens_input = target_tokenized["input_ids"][0][target_mask]
+    if hasattr(tokens_input, 'cpu'):
+        tokens_input = tokens_input.cpu().numpy()
+    target_token_str = target_tokenizer.convert_ids_to_tokens(tokens_input)
+
+    tokens_input = drug_tokenized["input_ids"][0][drug_mask]
+    if hasattr(tokens_input, 'cpu'):
+        tokens_input = tokens_input.cpu().numpy()
+    drug_token_str = drug_tokenizer.convert_ids_to_tokens(tokens_input)
+
+    # Extract subset of attention weights
+    if hasattr(crossattention_weights, 'cpu'):
+        crossattention_weights = crossattention_weights.cpu().numpy()
+    
+    subset = crossattention_weights[target_mask][:, drug_mask]
+    height, width = subset.shape
+    
+    fig, ax = plt.subplots(
+        figsize=(width * 0.2 + 2, height * 0.2 + 3),
+        dpi=300
+    )
+    im = ax.imshow(subset, 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()
+    for i in range(height):
+        for j in range(width):
+            val = subset[i, j]
+            if val > max_val / 2:
+                # Extract just the digits after the decimal (no leading '0.')
+                text = f"{val % 1:.2f}"[2:]
+                plt.text(j, i, text,
+                        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()
+    
+    return img
+
+
+def plot_presum(tokenized_input, affinities, scaler, w, b, target_tokenizer, 
+               raw_affinities=False):
+    """
+    Generates an annotated 1D heatmap of token-level contribution scores.
+
+    Args:
+        tokenized_input (dict): Output of a tokenizer with keys:
+            - 'input_ids' (torch.Tensor): token ID sequences, shape (1, seq_len)
+            - 'attention_mask' (torch.Tensor): mask indicating padding tokens
+        affinities (torch.Tensor): Final layer summation affinity contributions from the model, shape (1, seq_len)
+        scaler (object): Fitted scaler with `mean_` and `std_` attributes for inverse-transform.
+        w (float): Weight applied to the summed affinities before bias.
+        b (float): Bias added to the summed affinities.
+        target_tokenizer: Target tokenizer instance.
+        raw_affinities (bool): If True, plot raw (signed) contributions on a blue—white—red scale.
+            If False, enforce non-negative contributions and use a white—red scale.
+            Default: False
+
+    Returns:
+        PIL.Image: The generated contribution visualization image.
+        
+    Raises:
+        ValueError: If `sum(transformed_affinities) < 0` when `raw_affinities=False`.
+    """
+    colors = [
+        (1.0, 0.95, 0.95),  
+        (1.0, 0.5, 0.5),  
+        (0.8, 0.0, 0.0)   
+    ]
+
+    custom_reds = LinearSegmentedColormap.from_list("CustomReds", colors)
+    
+    # Convert tensors to numpy if needed
+    if hasattr(affinities, 'cpu'):
+        affinities = affinities.cpu().numpy()
+    if hasattr(w, 'cpu'):
+        w = w.cpu().numpy()
+    if hasattr(b, 'cpu'):
+        b = b.cpu().numpy()
+    
+    # Apply transformations
+    affinities = w * (affinities[0]) + b / len(affinities[0])
+    affinities = (affinities * scaler.std_) + scaler.mean_ / len(affinities)
+    
+    if sum(affinities) < 0 and not raw_affinities:
+        raise ValueError("Cannot use non-raw affinities with negative binding affinity prediction")
+
+    # Get token strings
+    tokens_input = tokenized_input["input_ids"][0]
+    if hasattr(tokens_input, 'cpu'):
+        tokens_input = tokens_input.cpu().numpy()
+    token_str = target_tokenizer.convert_ids_to_tokens(tokens_input)
+
+    # Handle padding
+    pad_mask = tokenized_input["attention_mask"][0] == 0
+    if hasattr(pad_mask, 'cpu'):
+        pad_mask = pad_mask.cpu().numpy()
+    
+    padding_affinities_sum = affinities[pad_mask].sum()
+    non_padding_affinities = affinities[~pad_mask]
+    processed_affinities = non_padding_affinities + padding_affinities_sum/len(non_padding_affinities)
+
+    # Make affinities non-negative if requested
+    if not raw_affinities:
+        all_negative_non_paddings = processed_affinities[processed_affinities < 0]
+
+        while(len(all_negative_non_paddings) > 0):
+            all_positive_non_paddings = processed_affinities[processed_affinities > 0]
+
+            processed_affinities[processed_affinities < 0] = 0
+            processed_affinities[processed_affinities > 0] = all_positive_non_paddings + all_negative_non_paddings.sum()/len(all_positive_non_paddings)
+            all_negative_non_paddings = processed_affinities[processed_affinities < 0]
+
+    # Create visualization
+    max_per_row = 20
+    n = len(processed_affinities)
+    n_rows = int(np.ceil(n / max_per_row))
+    grid = np.full((n_rows, max_per_row), np.nan)
+    grid.flat[:n] = processed_affinities
+
+    fig, ax = plt.subplots(
+        figsize = (max_per_row * 1, n_rows * 1 + 2),
+        dpi = 300
+    )
+
+    ax.set_xticks([])
+    ax.set_yticks([])
+
+    im = ax.imshow(
+        grid,
+        aspect='equal',
+        cmap='bwr' if raw_affinities else custom_reds,
+        vmin=np.nanmin(grid) if not raw_affinities else -max(abs(np.nanmin(grid)), abs(np.nanmax(grid))),
+        vmax=np.nanmax(grid) if not raw_affinities else max(abs(np.nanmin(grid)), abs(np.nanmax(grid))),
+    )
+
+    def wrap_text(text, width=8):
+        return '\n'.join(text[i:i+width] for i in range(0, len(text), width))
+
+    for idx, val in enumerate(processed_affinities):
+        r, c = divmod(idx, max_per_row)
+        wrapped_token = wrap_text(token_str[idx], width=8)
+        ax.text(c, r, f"{val:.2f}\n{wrapped_token}",
+                ha='center', va='center', fontsize=8)
+        
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('bottom', size=0.2, pad=0.3)
+    cbar = fig.colorbar(im, cax=cax, orientation='horizontal')
+    cbar.set_label("Contribution")
+    
+    # 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()
+    
+    return img

app.py

@@ -10,6 +10,9 @@ import os
 from pathlib import Path
 import logging
 
+# Import visualization functions
+from analysis import plot_crossattention_weights, plot_presum
+
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -93,9 +96,6 @@ class DrugTargetInteractionApp:
                 return_tensors="pt"
             ).to(self.device)
 
-            logger.info(f"Drug inputs: {drug_inputs}")
-            logger.info(f"Target inputs: {target_inputs}")
-
             # Make prediction
             with torch.no_grad():
                 prediction = self.model(target_inputs, drug_inputs)
@@ -111,58 +111,118 @@ class DrugTargetInteractionApp:
         except Exception as e:
             logger.error(f"Prediction error: {str(e)}")
             return f"Error during prediction: {str(e)}"
-    
-    def get_attention_visualization(self, target_sequence, drug_smiles, max_length=512):
-        """Get attention weights for visualization"""
+
+    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
+            
+        Returns:
+            tuple: (cross_attention_image, raw_contribution_image, normalized_contribution_image, status_message)
+        """
         if self.model is None:
-            return None, "Model not loaded"
+            return None, None, None, "Error: Model not loaded. Please load a model first."
         
         try:
-            # Enable interpretation mode
-            self.model.INTERPR_ENABLE_MODE()
-            
             # Tokenize inputs
             target_inputs = self.target_tokenizer(
                 target_sequence,
-                padding=True,
-                truncation=True,
-                max_length=max_length,
+                padding="max_length", 
+                truncation=True, 
+                max_length=512,
                 return_tensors="pt"
             ).to(self.device)
             
             drug_inputs = self.drug_tokenizer(
                 drug_smiles,
-                padding=True,
-                truncation=True,
-                max_length=max_length,
+                padding="max_length", 
+                truncation=True, 
+                max_length=512,
                 return_tensors="pt"
             ).to(self.device)
-            
-            # Make prediction to get attention weights
+
+            # Make prediction and extract visualization data
             with torch.no_grad():
-                _ = self.model(target_inputs, drug_inputs)
+                prediction = self.model(target_inputs, drug_inputs)
                 
-                # Get attention weights
-                attention_weights = self.model.model.crossattention_weights
-                if attention_weights is not None:
-                    # Convert to numpy for visualization
-                    attention_weights = attention_weights.cpu().numpy()
-                    
-                    # Get tokens for visualization
-                    target_tokens = self.target_tokenizer.convert_ids_to_tokens(
-                        target_inputs["input_ids"][0], skip_special_tokens=True
-                    )
-                    drug_tokens = self.drug_tokenizer.convert_ids_to_tokens(
-                        drug_inputs["input_ids"][0], skip_special_tokens=True
+                # Unscale if scaler exists
+                if self.model.scaler is not None:
+                    prediction = self.model.unscale(prediction)
+                
+                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[0][0]  # Shape: (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
+                
+            # Generate visualizations
+            try:
+                # 1. Cross-attention heatmap
+                cross_attention_img = plot_crossattention_weights(
+                    target_inputs["attention_mask"][0],
+                    drug_inputs["attention_mask"][0],
+                    target_inputs,
+                    drug_inputs,
+                    cross_attention_weights,
+                    self.target_tokenizer,
+                    self.drug_tokenizer
+                )
+            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 = plot_presum(
+                    target_inputs,
+                    presum_values,
+                    scaler,
+                    w,
+                    b,
+                    self.target_tokenizer,
+                    raw_affinities=False
+                )
+            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:
+                    raw_img = plot_presum(
+                        target_inputs,
+                        presum_values,
+                        scaler,
+                        w,
+                        b,
+                        self.target_tokenizer,
+                        raw_affinities=True
                     )
-                    
-                    return attention_weights, target_tokens, drug_tokens, "Attention visualization ready"
                 else:
-                    return None, None, None, "No attention weights available"
-                    
+                    logger.info("Skipping raw affinities visualization as pKd <= 0")
+            except Exception as e:
+                logger.error(f"Raw contribution visualization error: {str(e)}")
+                raw_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)"
+                
+            return cross_attention_img, raw_img, normalized_img, status_msg
+            
         except Exception as e:
-            logger.error(f"Attention visualization error: {str(e)}")
-            return None, None, None, f"Error: {str(e)}"
+            logger.error(f"Visualization error: {str(e)}")
+            return None, None, None, f"Error during visualization: {str(e)}"
+
 
 # Initialize the app
 app = DrugTargetInteractionApp()
@@ -174,6 +234,13 @@ def predict_wrapper(target_seq, drug_smiles):
     
     return app.predict_interaction(target_seq, drug_smiles)
 
+def visualize_wrapper(target_seq, drug_smiles):
+    """Wrapper function for visualization"""
+    if not target_seq.strip() or not drug_smiles.strip():
+        return None, None, None, "Please provide both target sequence and drug SMILES."
+
+    return app.visualize_interaction(target_seq, drug_smiles)
+
 def load_model_wrapper(model_path):
     """Wrapper function to load model"""
     if app.load_model(model_path):
@@ -210,7 +277,10 @@ with gr.Blocks(title="Drug-Target Interaction Predictor", theme=gr.themes.Soft()
                     lines=2
                 )
                 
-                predict_btn = gr.Button("🚀 Predict Interaction", variant="primary", size="lg")
+                # Buttons side by side
+                with gr.Row():
+                    predict_btn = gr.Button("🚀 Predict Interaction", variant="primary", size="lg")
+                    visualize_btn = gr.Button("📊 Visualize Interaction", variant="secondary", size="lg")
             
             with gr.Column(scale=1):
                 prediction_output = gr.Textbox(
@@ -219,6 +289,43 @@ with gr.Blocks(title="Drug-Target Interaction Predictor", theme=gr.themes.Soft()
                     lines=3
                 )
         
+        # Visualization outputs section
+        gr.HTML("<h3 style='margin-top: 30px; color: #2E86AB;'>📈 Interaction Visualizations</h3>")
+        
+        with gr.Row():
+            with gr.Column():
+                viz_image1 = gr.Image(
+                    label="Cross-Attention Heatmap",
+                    type="pil",
+                    interactive=False,
+                    container=True,
+                    height=300
+                )
+            
+            with gr.Column():
+                viz_image2 = gr.Image(
+                    label="Raw pKd Contribution Visualization",
+                    type="pil", 
+                    interactive=False,
+                    container=True,
+                    height=300
+                )
+            
+            with gr.Column():
+                viz_image3 = gr.Image(
+                    label="Normalized pKd Contribution Visualization",
+                    type="pil",
+                    interactive=False,
+                    container=True,
+                    height=300
+                )
+        
+        viz_status = gr.Textbox(
+            label="Visualization Status",
+            interactive=False,
+            lines=2
+        )
+        
         # Example inputs
         gr.HTML("<h3 style='margin-top: 20px; color: #2E86AB;'>📚 Example Inputs:</h3>")
         
@@ -239,11 +346,18 @@ with gr.Blocks(title="Drug-Target Interaction Predictor", theme=gr.themes.Soft()
             cache_examples=False
         )
         
+        # Button click events
         predict_btn.click(
             fn=predict_wrapper,
             inputs=[target_input, drug_input],
             outputs=prediction_output
         )
+        
+        visualize_btn.click(
+            fn=visualize_wrapper,
+            inputs=[target_input, drug_input],
+            outputs=[viz_image1, viz_image2, viz_image3, viz_status]
+        )
     
     with gr.Tab("⚙️ Model Settings"):
         gr.HTML("<h3 style='color: #2E86AB;'>Model Configuration</h3>")
@@ -286,14 +400,20 @@ with gr.Blocks(title="Drug-Target Interaction Predictor", theme=gr.themes.Soft()
         - Cross-attention for drug-target interaction modeling
         - Dropout for regularization
         - Layer normalization for stable training
-        - Interpretability mode for attention visualization
+        - Interpretability mode for contribution and attention visualization
         
         ### Usage Tips:
         1. Load your trained model using the Model Settings tab
         2. Enter a RNA sequence and drug SMILES
         3. Click "Predict Interaction" to get binding affinity prediction
+        4. Click "Visualize Interaction" to see detailed interaction analysis
         
         For best results, ensure your input sequences are properly formatted and within reasonable length limits.
+        
+        ### Visualization Features:
+        - **Cross-Attention Heatmap**: Shows cross-attention between drug and target tokens
+        - **Raw pKd Contribution**: Shows raw signed contributions (only when pKd > 0)
+        - **Normalized pKd Contribution**: Shows normalized non-negative contributions
         """)
 
 # Launch the app