Update app.py

app.py

@@ -1,17 +1,13 @@
 # app.py
 import gradio as gr
 import torch
-from transformers import AutoModelForMaskedLM, AutoTokenizer, pipeline, RobertaModel, RobertaTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForMaskedLM, AutoTokenizer, pipeline, BitsAndBytesConfig
 from rdkit import Chem
 from rdkit.Chem import Draw, rdFMCS
 from rdkit.Chem.Draw import MolToImage
 # 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 # 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
@@ -58,14 +54,13 @@ def load_optimized_models():
 
     logger.info(f"Loading models on device: {device} with dtype: {torch_dtype}")
 
-    # Model names
+    # Model name
     model_name = "seyonec/PubChem10M_SMILES_BPE_450k"
 
-    # Load tokenizers (these don't need quantization)
+    # Load tokenizer (doesn't need quantization)
     fill_mask_tokenizer = AutoTokenizer.from_pretrained(model_name)
-    attention_tokenizer = RobertaTokenizer.from_pretrained(model_name)
 
-    # Load models with quantization if available
+    # Load model with quantization if available
     model_kwargs = {
         "torch_dtype": torch_dtype,
     }
@@ -85,35 +80,21 @@ def load_optimized_models():
             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()
+        fill_mask_model.eval() # Set model to evaluation mode for inference
 
         # 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
+        return fill_mask_tokenizer, fill_mask_model, fill_mask_pipeline
 
     except Exception as e:
         logger.error(f"Error loading optimized models: {e}")
@@ -129,17 +110,13 @@ def load_standard_models(model_name):
     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
+    return fill_mask_tokenizer, fill_mask_model, fill_mask_pipeline
 
 # Load models with optimizations
-fill_mask_tokenizer, fill_mask_model, fill_mask_pipeline, attention_model, attention_tokenizer = load_optimized_models()
+fill_mask_tokenizer, fill_mask_model, fill_mask_pipeline = load_optimized_models()
 
 # --- Memory Management Utilities ---
 def clear_gpu_cache():
@@ -249,57 +226,6 @@ def predict_and_visualize_masked_smiles(smiles_mask, substructure_smarts_highlig
     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 neuron-by-neuron attention view as HTML.
-    Optimized with memory management and mixed precision.
-    """
-    if not sentence_a or not sentence_b:
-        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']
-
-        # 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():
-            # 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)
-
-        # 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 # .data should provide the HTML string
-
-        # Add D3 and jQuery CDN links to the HTML string for better rendering in Gradio
-        html_with_deps = f"""
-        <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:
-        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):
     """
     Displays a 2D image of a molecule from its SMILES string.
@@ -346,26 +272,6 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
             outputs=[predictions_table, img_out_1, img_out_2, img_out_3, img_out_4, img_out_5, status_masked]
         )
 
-    with gr.Tab("Attention Visualization"):
-        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 Neuron View") # Changed label for clarity
-
-        # Automatically populate on load for the default example
-        demo.load(
-            lambda: visualize_attention_bertviz("CCCCC[C@@H](Br)CC", "CCCCC[C@H](Br)CC"),
-            inputs=None,
-            outputs=[attention_html_output]
-        )
-        visualize_button_attn.click(
-            visualize_attention_bertviz,
-            inputs=[smiles_a_input_attn, smiles_b_input_attn],
-            outputs=[attention_html_output]
-        )
-
     with gr.Tab("Molecule Viewer"):
         gr.Markdown("Enter a SMILES string to display its 2D structure.")
         smiles_input_viewer = gr.Textbox(label="SMILES String", value="C1=CC=CC=C1")
@@ -386,4 +292,4 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
         )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()