Update app.py

app.py

@@ -2,6 +2,7 @@ from rdkit import Chem
 from rdkit.Chem import Draw
 from transformers import pipeline
 import gradio as gr
+from interpretability import save_high_quality_png
 
 model_checkpoint = "FartLabs/FART_Augmented"
 classifier = pipeline("text-classification", model=model_checkpoint, top_k=None)
@@ -17,9 +18,9 @@ def process_smiles(smiles):
     predictions = classifier(canonical_smiles)
     
     # Generate molecule image
-    img_path = "molecule.png"
-    img = Draw.MolToImage(mol)
-    img.save(img_path)
+    img_path = "molecule"
+    filepath= "molecule.png"
+    save_high_quality_png(smiles, molecule)
     
     # Convert predictions to a friendly format
     prediction_dict = {pred["label"]: pred["score"] for pred in predictions[0]}