Update app.py

app.py

@@ -2,12 +2,74 @@ from rdkit import Chem
 from rdkit.Chem import Draw
 from transformers import pipeline
 import gradio as gr
-from interpretability import save_high_quality_png
+from rdkit import Chem
+from rdkit.Chem import Draw
+from rdkit.Chem.Draw import SimilarityMaps
+import io
+from PIL import Image
+import numpy as np
+import rdkit
+
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from transformers_interpret import SequenceClassificationExplainer
+
+model_name = "FartLabs/FART_Augmented"
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+cls_explainer = SequenceClassificationExplainer(model, tokenizer)
+
+def save_high_quality_png(smiles, title, bw=True, padding=0.05):
+    """
+    Generates a high-quality PNG of atom-wise gradients or importance scores for a molecule.
+    Parameters:
+    - smiles (str): The SMILES string of the molecule to visualize.
+    - token_importance (list): List of importance scores for each atom.
+    - bw (bool): If True, renders the molecule in black and white.
+    - padding (float): Padding for molecule drawing.
+    - output_file (str): Path to save the high-quality PNG file.
+    Returns:
+    - None
+    """
+    
+    # Convert SMILES string to RDKit molecule object
+    molecule = Chem.MolFromSmiles(smiles)
+    Chem.rdDepictor.Compute2DCoords(molecule)
+    
+    # Get token importance scores and map to atoms
+    token_importance = cls_explainer(smiles)
+    atom_importance = [c[1] for c in token_importance if c[0].isalpha()]
+    num_atoms = molecule.GetNumAtoms()
+    atom_importance = atom_importance[:num_atoms]
+    
+    # Set a large canvas size for high resolution
+    d = Draw.MolDraw2DCairo(1500, 1500)
+
+    dopts = d.drawOptions()
+    dopts.padding = padding    
+    dopts.maxFontSize = 2000
+    dopts.bondLineWidth = 5
+
+    # Optionally set black and white palette
+    if bw:
+        d.drawOptions().useBWAtomPalette()
+    
+    # Generate and display a similarity map based on atom importance scores
+    SimilarityMaps.GetSimilarityMapFromWeights(molecule, atom_importance, draw2d=d)
 
+    # Draw molecule with color highlights
+    d.FinishDrawing()
+    
+    # Save to PNG file with high quality
+    with open(f"{title}.png", "wb") as png_file:
+        png_file.write(d.GetDrawingText())
+    
+    return None
+    
 model_checkpoint = "FartLabs/FART_Augmented"
 classifier = pipeline("text-classification", model=model_checkpoint, top_k=None)
 
-def process_smiles(smiles):
+def process_smiles(smiles, compute_explanation):
     # Validate and canonicalize SMILES
     mol = Chem.MolFromSmiles(smiles)
     if mol is None:
@@ -18,9 +80,14 @@ def process_smiles(smiles):
     predictions = classifier(canonical_smiles)
     
     # Generate molecule image
-    img_path = "molecule"
-    filepath= "molecule.png"
-    save_high_quality_png(smiles, img_path)
+    if compute_explanation:
+        img_path = "molecule"
+        filepath= "molecule.png"
+        save_high_quality_png(smiles, img_path)
+    else:
+        filepath = "molecule.png"
+        img = Draw.MolToImage(mol)
+        img.save(filepath)
     
     # Convert predictions to a friendly format
     prediction_dict = {pred["label"]: pred["score"] for pred in predictions[0]}
@@ -29,7 +96,10 @@ def process_smiles(smiles):
 
 iface = gr.Interface(
     fn=process_smiles,
-    inputs=gr.Textbox(label="Input SMILES", value="O1[C@H](CO)[C@@H](O)[C@H](O)[C@@H](O)[C@H]1O[C@@]2(O[C@@H]([C@@H](O)[C@@H]2O)CO)CO"),
+    inputs=[
+        gr.Textbox(label="Input SMILES", value="O1[C@H](CO)[C@@H](O)[C@H](O)[C@@H](O)[C@H]1O[C@@]2(O[C@@H]([C@@H](O)[C@@H]2O)CO)CO"),
+        gr.Checkbox(label="Compute Explanation (Takes 60s)", value=False),
+    ],
     outputs=[
         gr.Label(num_top_classes=3, label="Classification Probabilities"),
         gr.Image(type="filepath", label="Molecule Image"),