Update app.py

app.py

@@ -1,28 +1,151 @@
 # app.py
 import gradio as gr
 import torch
-from transformers import AutoModelForMaskedLM, AutoTokenizer, pipeline, RobertaModel, RobertaTokenizer
+from transformers import AutoModelForMaskedLM, AutoTokenizer, pipeline, RobertaModel, RobertaTokenizer, BitsAndBytesConfig
 from rdkit import Chem
 from rdkit.Chem import Draw, rdFMCS
 from rdkit.Chem.Draw import MolToImage
-from PIL import Image  # Corrected Line
+# PIL is imported as Image by rdkit.Chem.Draw.MolToImage, but explicit import is good practice if used directly.
+# from PIL import Image
 import pandas as pd
-from bertviz import head_view
-from IPython.core.display import HTML
+from bertviz import head_view # For potential future use or if other parts rely on it
+from bertviz import neuron_view as neuron_view_function # Specific import for neuron_view function
+# IPython.core.display.HTML is generally for notebooks. Gradio's gr.HTML handles HTML strings directly.
+# from IPython.core.display import HTML
 import io
 import base64
+import logging
 
-# --- Model and Tokenizer Loading ---
-# Masked LM Model
-fill_mask_model_name = "seyonec/PubChem10M_SMILES_BPE_450k"
-fill_mask_tokenizer = AutoTokenizer.from_pretrained(fill_mask_model_name)
-fill_mask_model = AutoModelForMaskedLM.from_pretrained(fill_mask_model_name)
-fill_mask_pipeline = pipeline('fill-mask', model=fill_mask_model, tokenizer=fill_mask_tokenizer)
+# Set up logging to monitor quantization effects
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# Roberta Model for Attention
-attention_model_name = 'seyonec/PubChem10M_SMILES_BPE_450k' # Can be same or different as needed
-attention_model = RobertaModel.from_pretrained(attention_model_name, output_attentions=True)
-attention_tokenizer = RobertaTokenizer.from_pretrained(attention_model_name)
+# --- Quantization Configuration ---
+def get_quantization_config():
+    """
+    Configure 8-bit quantization for model optimization.
+    Falls back gracefully if bitsandbytes is not available.
+    """
+    try:
+        # 8-bit quantization configuration - good balance of speed and quality
+        quantization_config = BitsAndBytesConfig(
+            load_in_8bit=True,
+            bnb_8bit_compute_dtype=torch.float16,
+            bnb_8bit_use_double_quant=True,  # Nested quantization for better compression
+        )
+        logger.info("8-bit quantization configuration loaded successfully")
+        return quantization_config
+    except ImportError:
+        logger.warning("bitsandbytes not available, falling back to standard loading")
+        return None
+    except Exception as e:
+        logger.warning(f"Quantization setup failed: {e}, using standard loading")
+        return None
+
+def get_torch_dtype():
+    """Get appropriate torch dtype based on available hardware."""
+    if torch.cuda.is_available():
+        return torch.float16  # Use half precision on GPU
+    else:
+        return torch.float32  # Keep full precision on CPU
+
+# --- Optimized Model Loading ---
+def load_optimized_models():
+    """Load models with quantization and other optimizations."""
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    torch_dtype = get_torch_dtype()
+    quantization_config = get_quantization_config()
+
+    logger.info(f"Loading models on device: {device} with dtype: {torch_dtype}")
+
+    # Model names
+    model_name = "seyonec/PubChem10M_SMILES_BPE_450k"
+
+    # Load tokenizers (these don't need quantization)
+    fill_mask_tokenizer = AutoTokenizer.from_pretrained(model_name)
+    attention_tokenizer = RobertaTokenizer.from_pretrained(model_name)
+
+    # Load models with quantization if available
+    model_kwargs = {
+        "torch_dtype": torch_dtype,
+    }
+
+    if quantization_config is not None and torch.cuda.is_available(): # Quantization typically for GPU
+        model_kwargs["quantization_config"] = quantization_config
+        # device_map="auto" is often used with bitsandbytes for automatic distribution
+        model_kwargs["device_map"] = "auto"
+    elif torch.cuda.is_available():
+        model_kwargs["device_map"] = "auto" # For non-quantized GPU loading
+    else:
+        model_kwargs["device_map"] = None # For CPU
+
+    try:
+        # Masked LM Model
+        fill_mask_model = AutoModelForMaskedLM.from_pretrained(
+            model_name,
+            **model_kwargs
+        )
+
+        # RoBERTa model for attention
+        attention_model_kwargs = model_kwargs.copy()
+        attention_model_kwargs["output_attentions"] = True
+
+        attention_model = RobertaModel.from_pretrained(
+            model_name,
+            **attention_model_kwargs
+        )
+
+        # Set models to evaluation mode for inference
+        fill_mask_model.eval()
+        attention_model.eval()
+
+        # Create optimized pipeline
+        # Let pipeline infer device from model if possible, or set based on model's device
+        pipeline_device = fill_mask_model.device.index if hasattr(fill_mask_model.device, 'type') and fill_mask_model.device.type == "cuda" else -1
+
+
+        fill_mask_pipeline = pipeline(
+            'fill-mask',
+            model=fill_mask_model,
+            tokenizer=fill_mask_tokenizer,
+            device=pipeline_device, # Use model's device
+            # torch_dtype=torch_dtype # Pipeline might infer this or it might conflict
+        )
+
+        logger.info("Models loaded successfully with optimizations")
+        return fill_mask_tokenizer, fill_mask_model, fill_mask_pipeline, attention_model, attention_tokenizer
+
+    except Exception as e:
+        logger.error(f"Error loading optimized models: {e}")
+        # Fallback to standard loading
+        logger.info("Falling back to standard model loading...")
+        return load_standard_models(model_name)
+
+def load_standard_models(model_name):
+    """Fallback standard model loading without quantization."""
+    fill_mask_tokenizer = AutoTokenizer.from_pretrained(model_name)
+    fill_mask_model = AutoModelForMaskedLM.from_pretrained(model_name)
+    # Determine device for standard loading
+    device_idx = 0 if torch.cuda.is_available() else -1
+    fill_mask_pipeline = pipeline('fill-mask', model=fill_mask_model, tokenizer=fill_mask_tokenizer, device=device_idx)
+
+    attention_model = RobertaModel.from_pretrained(model_name, output_attentions=True)
+    attention_tokenizer = RobertaTokenizer.from_pretrained(model_name)
+
+    if torch.cuda.is_available():
+        fill_mask_model.to("cuda")
+        attention_model.to("cuda")
+
+    return fill_mask_tokenizer, fill_mask_model, fill_mask_pipeline, attention_model, attention_tokenizer
+
+# Load models with optimizations
+fill_mask_tokenizer, fill_mask_model, fill_mask_pipeline, attention_model, attention_tokenizer = load_optimized_models()
+
+# --- Memory Management Utilities ---
+def clear_gpu_cache():
+    """Clear CUDA cache to free up memory."""
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
 
 # --- Helper Functions from Notebook (adapted) ---
 def get_mol(smiles):
@@ -51,24 +174,40 @@ def get_image_with_highlight(mol, atomset=None, size=(300, 300)):
     if mol is None:
         return None
     highlight_color = (0, 1, 0, 0.5) # Green with some transparency
+
+    # Ensure atomset contains integers if not None or empty
+    valid_atomset = []
+    if atomset:
+        try:
+            valid_atomset = [int(a) for a in atomset]
+        except ValueError:
+            logger.warning(f"Invalid atom in atomset: {atomset}. Proceeding without highlighting problematic atoms.")
+            valid_atomset = [int(a) for a in atomset if str(a).isdigit()] # Filter out non-integers
+
     img = MolToImage(mol, size=size, fitImage=True,
-                     highlightAtoms=atomset if atomset else [],
-                     highlightAtomColors={i: highlight_color for i in atomset} if atomset else {})
+                     highlightAtoms=valid_atomset if valid_atomset else [],
+                     highlightAtomColors={i: highlight_color for i in valid_atomset} if valid_atomset else {})
     return img
 
-# --- Gradio Interface Functions ---
+# --- Optimized Gradio Interface Functions ---
 
 def predict_and_visualize_masked_smiles(smiles_mask, substructure_smarts_highlight="CC=CC"):
     """
     Predicts masked tokens in a SMILES string, shows scores, and visualizes molecules.
+    Optimized with memory management. Returns 7 items for Gradio outputs.
     """
     if fill_mask_tokenizer.mask_token not in smiles_mask:
-        return pd.DataFrame(), [None]*5, "Error: Input SMILES must contain a mask token (e.g., <mask>)."
+        # Return 7 items for the 7 output components
+        return pd.DataFrame(), None, None, None, None, None, "Error: Input SMILES must contain a mask token (e.g., <mask>)."
 
     try:
-        predictions = fill_mask_pipeline(smiles_mask, top_k=10) # Get more to filter for valid ones
+        # Use torch.no_grad() for inference to save memory
+        with torch.no_grad():
+            predictions = fill_mask_pipeline(smiles_mask, top_k=10) # Get more to filter for valid ones
     except Exception as e:
-        return pd.DataFrame(), [None]*5, f"Error during prediction: {str(e)}"
+        clear_gpu_cache()  # Clear cache on error
+        # Return 7 items
+        return pd.DataFrame(), None, None, None, None, None, f"Error during prediction: {str(e)}"
 
     results_data = []
     image_list = []
@@ -85,13 +224,13 @@ def predict_and_visualize_masked_smiles(smiles_mask, substructure_smarts_highlig
         if mol:
             results_data.append({"Predicted SMILES": predicted_smiles, "Score": f"{score:.4f}"})
 
-            atom_matches = []
+            atom_matches_indices = []
             if substructure_smarts_highlight:
                 matches = find_matches_one(mol, substructure_smarts_highlight)
                 if matches:
-                    atom_matches = list(matches[0]) # Highlight first match
+                    atom_matches_indices = list(matches[0]) # Highlight first match
 
-            img = get_image_with_highlight(mol, atomset=atom_matches)
+            img = get_image_with_highlight(mol, atomset=atom_matches_indices)
             image_list.append(img)
             valid_predictions_count += 1
 
@@ -100,50 +239,66 @@ def predict_and_visualize_masked_smiles(smiles_mask, substructure_smarts_highlig
         image_list.append(None)
 
     df_results = pd.DataFrame(results_data)
-    return df_results, image_list, "Prediction successful." if valid_predictions_count > 0 else "No valid molecules found for top predictions."
+
+    # Clear cache after inference
+    clear_gpu_cache()
+
+    status_message = "Prediction successful." if valid_predictions_count > 0 else "No valid molecules found for top predictions."
+
+    # Unpack image_list into individual image outputs + df_results + status_message
+    return df_results, image_list[0], image_list[1], image_list[2], image_list[3], image_list[4], status_message
 
 
 def visualize_attention_bertviz(sentence_a, sentence_b):
     """
-    Generates and displays BertViz attention head view as HTML.
+    Generates and displays BertViz neuron-by-neuron attention view as HTML.
+    Optimized with memory management and mixed precision.
     """
     if not sentence_a or not sentence_b:
-        return "Please provide two SMILES strings."
+        return "<p style='color:red;'>Please provide two SMILES strings.</p>"
     try:
         inputs = attention_tokenizer.encode_plus(sentence_a, sentence_b, return_tensors='pt', add_special_tokens=True)
         input_ids = inputs['input_ids']
 
-        # Ensure model is in eval mode and no_grad for inference
+        # Move to appropriate device if using GPU
+        if torch.cuda.is_available() and hasattr(attention_model, 'device'):
+            input_ids = input_ids.to(attention_model.device)
+
+        # Ensure model is in eval mode and use no_grad for inference
         attention_model.eval()
         with torch.no_grad():
-            attention_outputs = attention_model(input_ids)
+            # Use autocast for mixed precision if on CUDA
+            if torch.cuda.is_available() and hasattr(torch.cuda, 'amp') and hasattr(torch.cuda.amp, 'autocast'): # Check for amp
+                with torch.cuda.amp.autocast(dtype=torch.float16 if get_torch_dtype() == torch.float16 else None):
+                    attention_outputs = attention_model(input_ids)
+            else:
+                attention_outputs = attention_model(input_ids)
 
         attention = attention_outputs[-1] # Last item in the tuple is attentions
         input_id_list = input_ids[0].tolist()
         tokens = attention_tokenizer.convert_ids_to_tokens(input_id_list)
 
-        html_object = head_view(attention, tokens, display_mode="light") # Use light mode for better Gradio compatibility
+        # Using the specifically imported neuron_view_function
+        html_object = neuron_view_function(attention, tokens)
 
         # Extract HTML string from the IPython.core.display.HTML object
-        html_string = html_object.data
-
-        # Embed JavaScript directly if needed, or ensure Gradio's HTML component handles it.
-        # BertViz often requires D3.js and jQuery. Gradio's HTML component might not execute all JS.
-        # For robustness, it's better if head_view produces self-contained HTML or if Gradio supports JS execution.
-        # A common workaround is to serve the HTML and use an iframe, or save to file and link.
-        # Here, we'll return the raw HTML string and let Gradio's gr.HTML handle it.
+        html_string = html_object.data # .data should provide the HTML string
 
         # Add D3 and jQuery CDN links to the HTML string for better rendering in Gradio
-        # This is a common workaround if Gradio's HTML component doesn't include these by default
-        # Note: This might still have limitations depending on Gradio's sandboxing.
         html_with_deps = f"""
-        <script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/2.0.0/jquery.min.js"></script>
-        <script src="https://cdnjs.cloudflare.com/ajax/libs/d3/3.5.8/d3.min.js"></script>
+        <script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/3.5.1/jquery.min.js"></script>
+        <script src="https://cdnjs.cloudflare.com/ajax/libs/d3/5.16.0/d3.min.js"></script>
         {html_string}
         """
+
+        # Clear cache after attention computation
+        clear_gpu_cache()
+
         return html_with_deps
     except Exception as e:
-        return f"Error generating attention visualization: {str(e)}"
+        clear_gpu_cache()  # Clear cache on error
+        logger.error(f"Error in visualize_attention_bertviz: {e}", exc_info=True)
+        return f"<p style='color:red;'>Error generating attention visualization: {str(e)}</p>"
 
 def display_molecule_image(smiles_string):
     """
@@ -192,12 +347,12 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
         )
 
     with gr.Tab("Attention Visualization"):
-        gr.Markdown("Enter two SMILES strings to visualize attention between them using BertViz. This may take a moment to render.")
+        gr.Markdown("Enter two SMILES strings to visualize **neuron-by-neuron attention** between them using BertViz. This may take a moment to render.")
         with gr.Row():
             smiles_a_input_attn = gr.Textbox(label="SMILES String A", value="CCCCC[C@@H](Br)CC")
             smiles_b_input_attn = gr.Textbox(label="SMILES String B", value="CCCCC[C@H](Br)CC")
         visualize_button_attn = gr.Button("Visualize Attention")
-        attention_html_output = gr.HTML(label="Attention Head View")
+        attention_html_output = gr.HTML(label="Attention Neuron View") # Changed label for clarity
 
         # Automatically populate on load for the default example
         demo.load(