app.py

import gradio as gr
import torch
import numpy as np
from transformers import AutoModel, AutoTokenizer, AutoConfig, RobertaModel
from modeling_dlmberta import InteractionModelATTNForRegression, StdScaler
from configuration_dlmberta import InteractionModelATTNConfig
from chemberta import ChembertaTokenizer
import json
import os
from pathlib import Path
import logging

# Import visualization functions
from analysis import plot_crossattention_weights, plot_presum
from PIL import Image, ImageDraw, ImageFont

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def create_placeholder_image(width=400, height=300, text="No visualization available", bg_color=(0, 0, 0, 0)):
    """
    Create a transparent placeholder image with text
    
    Args:
        width (int): Image width
        height (int): Image height  
        text (str): Text to display
        bg_color (tuple): Background color (R, G, B, A) - (0,0,0,0) for transparent
        
    Returns:
        PIL.Image: Transparent placeholder image
    """
    # Create image with transparent background
    img = Image.new('RGBA', (width, height), bg_color)
    draw = ImageDraw.Draw(img)
    
    # Try to use a default font, fallback to default if not available
    try:
        font = ImageFont.truetype("arial.ttf", 16)
    except:
        try:
            font = ImageFont.load_default()
        except:
            font = None
    
    # Get text size and position for centering
    if font:
        bbox = draw.textbbox((0, 0), text, font=font)
        text_width = bbox[2] - bbox[0]
        text_height = bbox[3] - bbox[1]
    else:
        # Rough estimation if no font available
        text_width = len(text) * 8
        text_height = 16
    
    x = (width - text_width) // 2
    y = (height - text_height) // 2
    
    # Draw text in gray
    draw.text((x, y), text, fill=(128, 128, 128, 255), font=font)
    
    return img

class DrugTargetInteractionApp:
    def __init__(self):
        self.model = None
        self.target_tokenizer = None
        self.drug_tokenizer = None
        self.scaler = None
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
    def load_model(self, model_path="./"):
        """Load the pre-trained model and tokenizers"""
        try:
            # Load configuration
            config = InteractionModelATTNConfig.from_pretrained(model_path)
            
            # Load drug encoder (ChemBERTa)
            drug_encoder_config = AutoConfig.from_pretrained("DeepChem/ChemBERTa-77M-MTR")
            drug_encoder_config.pooler = None
            drug_encoder = RobertaModel(config=drug_encoder_config, add_pooling_layer=False)
            
            # Load target encoder
            target_encoder = AutoModel.from_pretrained("IlPakoZ/RNA-BERTa9700")
            
            # Load scaler if exists
            scaler_path = os.path.join(model_path, "scaler.config")
            scaler = None
            if os.path.exists(scaler_path):
                scaler = StdScaler()
                scaler.load(model_path)
            
            self.model = InteractionModelATTNForRegression.from_pretrained(
                        model_path,
                        config=config,
                        target_encoder=target_encoder,
                        drug_encoder=drug_encoder, 
                        scaler=scaler
                    )

            self.model.to(self.device)
            self.model.eval()
            
            # Load tokenizers
            self.target_tokenizer = AutoTokenizer.from_pretrained(
                os.path.join(model_path, "target_tokenizer")
            )
            
            # Load drug tokenizer (ChemBERTa)
            vocab_file = os.path.join(model_path, "drug_tokenizer", "vocab.json")
            self.drug_tokenizer = ChembertaTokenizer(vocab_file)
            
            logger.info("Model and tokenizers loaded successfully!")
            return True
            
        except Exception as e:
            logger.error(f"Error loading model: {str(e)}")
            return False
    
    def predict_interaction(self, target_sequence, drug_smiles, max_length=512):
        """Predict drug-target interaction"""
        if self.model is None:
            return "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)

            # Make prediction
            self.model.INTERPR_DISABLE_MODE()
            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)
                
                prediction_value = prediction.cpu().numpy()[0][0]
            
            return f"Predicted Binding Affinity: {prediction_value:.4f}"
            
        except Exception as e:
            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
            
        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)
                
                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'}")
                    
                    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
            
            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
            
            # 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)"
                )
                
            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
app = DrugTargetInteractionApp()

def predict_wrapper(target_seq, drug_smiles):
    """Wrapper function for Gradio interface"""
    if not target_seq.strip() or not drug_smiles.strip():
        return "Please provide both target sequence and 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):
        return "Model loaded successfully!"
    else:
        return "Failed to load model. Check the path and files."

# Create Gradio interface
with gr.Blocks(title="Drug-Target Interaction Predictor", theme=gr.themes.Soft()) as demo:
    gr.HTML("""
    <div style="text-align: center; margin-bottom: 30px;">
        <h1 style="color: #2E86AB; font-size: 2.5em; margin-bottom: 10px;">
            🧬 Drug-Target Interaction Predictor
        </h1>
        <p style="font-size: 1.2em; color: #666;">
            Predict binding affinity between drugs and target RNA sequences using deep learning
        </p>
    </div>
    """)
    
    with gr.Tab("🔮 Prediction"):
        with gr.Row():
            with gr.Column(scale=1):
                target_input = gr.Textbox(
                    label="Target RNA Sequence",
                    placeholder="Enter RNA sequence (e.g., AUGCUAGCUAGUACGUA...)",
                    lines=4,
                    max_lines=6
                )
                
                drug_input = gr.Textbox(
                    label="Drug SMILES",
                    placeholder="Enter SMILES notation (e.g., CC(C)CC1=CC=C(C=C1)C(C)C(=O)O)",
                    lines=2
                )
                
                # 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(
                    label="Prediction Result",
                    interactive=False,
                    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,
                    value=create_placeholder_image(text="Cross-Attention Heatmap\n(Click Visualize to generate)")
                )
            
            with gr.Column():
                viz_image2 = gr.Image(
                    label="Raw pKd Contribution Visualization",
                    type="pil", 
                    interactive=False,
                    container=True,
                    height=300,
                    value=create_placeholder_image(text="Raw pKd Contribution\n(Click Visualize to generate)")
                )
            
            with gr.Column():
                viz_image3 = gr.Image(
                    label="Normalized pKd Contribution Visualization",
                    type="pil",
                    interactive=False,
                    container=True,
                    height=300,
                    value=create_placeholder_image(text="Normalized pKd Contribution\n(Click Visualize to generate)")
                )
        
        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>")
        
        examples = gr.Examples(
            examples=[
                [
                    "AUGCUAGCUAGUACGUAUAUCUGCACUGC",
                    "CC(C)CC1=CC=C(C=C1)C(C)C(=O)O"
                ],
                [
                    "AUGCGAUCGACGUACGUUAGCCGUAGCGUAGCUAGUGUAGCUAGUAGCU",
                    "C1=CC=C(C=C1)NC(=O)C2=CC=CC=N2"
                ]
            ],
            inputs=[target_input, drug_input],
            outputs=prediction_output,
            fn=predict_wrapper,
            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>")
        
        model_path_input = gr.Textbox(
            label="Model Path",
            value="./",
            placeholder="Path to model directory"
        )
        
        load_model_btn = gr.Button("🔥 Load Model", variant="secondary")
        model_status = gr.Textbox(
            label="Status",
            interactive=False,
            value="No model loaded"
        )
        
        load_model_btn.click(
            fn=load_model_wrapper,
            inputs=model_path_input,
            outputs=model_status
        )
    
    with gr.Tab("ℹ️ About"):
        gr.Markdown("""
        ## About This Application
        
        This application uses a deep learning model for predicting drug-target interactions. The model architecture includes:
        
        - **Target Encoder**: Processes RNA sequences
        - **Drug Encoder**: Processes molecular SMILES notation  
        - **Cross-Attention Mechanism**: Captures interactions between drugs and targets
        - **Regression Head**: Predicts binding affinity scores
        
        ### Input Requirements:
        - **Target Sequence**: RNA sequence of the target
        - **Drug SMILES**: Simplified Molecular Input Line Entry System notation
        
        ### Model Features:
        - Cross-attention for drug-target interaction modeling
        - Dropout for regularization
        - Layer normalization for stable training
        - 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
if __name__ == "__main__":
    # Try to load model on startup
    if os.path.exists("./config.json"):
        app.load_model("./")
    
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )