Update app.py

app.py

@@ -1,6 +1,8 @@
 import streamlit as st
 from trainer import Trainer
 import random
+from rdkit.Chem import Draw
+from rdkit import Chem
 
 class DrugGENConfig:
     submodel='CrossLoss'
@@ -124,8 +126,8 @@ if submitted:
         results = trainer.inference()
     st.success(f"Inference of {model_name} took {results['runtime']:.2f} seconds.")
 
-    with st.expander("Expand to see scores"):
-
+    with st.expander("Expand to see the generation performance scores"):
+        st.write("### Generation performance scores (novelty is calculated in comparison to the training dataset)")
         st.success(f"Validity: {results['fraction_valid']}")
         st.success(f"Uniqueness: {results['uniqueness']}")
         st.success(f"Novelty: {results['novelty']}")
@@ -134,6 +136,28 @@ if submitted:
         inference_drugs = f.read()
     st.download_button(label="Click to download generated molecules", data=inference_drugs, file_name=f'{model_name}_inference.smi', mime="text/plain")
 
+    
+    st.write("Structures of randomly selected 12 de novo molecules from the inference set:")
+    # from rdkit.Chem import Draw
+# img = Draw.MolsToGridImage(mol_list, molsPerRow=5, subImgSize=(250, 250), maxMols=num_mols,
+                           # legends=None, useSVG=True)
+    generated_molecule_list = inference_drugs.split("\n")
+
+    selected_molecules = random.choices(generated_molecule_list,k=12)
+
+    selected_molecules = [Chem.MolFromSmiles(mol) for mol in selected_molecules]
+
+    molecule_image = Draw.MolsToGridImage(
+        selected_molecules,
+        molsPerRow=3,
+        subImgSize=(250, 250),
+        maxMols=len(selected_molecules),
+        # legends=None,
+        useSVG=True
+    )
+    st.image(molecule_image)
+    
+
 else:
     st.warning("Please select a model to make inference")