Fix

app.py

@@ -47,6 +47,9 @@ base_dir = os.path.dirname(__file__)
 print("Base Dir : ", base_dir)
 
 import models.fm4m as fm4m
+import re
+from rdkit import RDLogger
+RDLogger.logger().setLevel(RDLogger.ERROR)
 
 
 # Function to display molecule image from SMILES
@@ -182,7 +185,7 @@ def generate(latent_vector, mask):
     selfies = gen_tokenizer.batch_decode(decoder_output, skip_special_tokens=True)
     outs = []
     for i in selfies:
-        outs.append(sf.decoder(i.replace("] [", "][")))
+        outs.append(sf.decoder(re.sub(r'\]\s*(.*?)\s*\[', r']\1[', i)))
     return outs
 
 
@@ -217,8 +220,22 @@ def generate_canonical(smiles):
         noise = i / 10
         perturbed_latent = perturb_latent(latent_vec, noise_scale=noise)
         gen = generate(perturbed_latent, mask)
-        gen_mol = Chem.MolToSmiles(Chem.MolFromSmiles(gen[0]))
-        if gen_mol != Chem.MolToSmiles(Chem.MolFromSmiles(smiles)): break
+        mol = Chem.MolFromSmiles(gen[0])
+        if mol:
+            gen_mol = Chem.MolToSmiles(mol)
+            if gen_mol != Chem.MolToSmiles(Chem.MolFromSmiles(smiles)): break
+        else:
+            print('Abnormal molecule:', gen[0])
+            gen_mols = []
+            for sel in gen[0].split('.'):
+                mol = Chem.MolFromSmiles(sel)
+                if mol:
+                    mol = Chem.MolToSmiles(mol)
+                    if mol != Chem.MolToSmiles(Chem.MolFromSmiles(smiles)):
+                        gen_mols.append(mol)
+            if len(gen_mols) > 0:
+                gen_mol = '.'.join(gen_mols)
+                break
 
     if gen_mol:
         # Calculate properties for ref and gen molecules