Update tool/comget/utils.py

tool/comget/utils.py

@@ -3,12 +3,29 @@ import numpy as np
 import torch
 import torch.nn as nn
 from torch.nn import functional as F
-from .moses.utils import get_mol
 from rdkit import Chem
    
 import numpy as np
 import threading
 
+
+def get_mol(smiles_or_mol):
+    '''
+    Loads SMILES/molecule into RDKit's object
+    '''
+    if isinstance(smiles_or_mol, str):
+        if len(smiles_or_mol) == 0:
+            return None
+        mol = Chem.MolFromSmiles(smiles_or_mol)
+        if mol is None:
+            return None
+        try:
+            Chem.SanitizeMol(mol)
+        except ValueError:
+            return None
+        return mol
+    return smiles_or_mol
+
 def set_seed(seed):
     random.seed(seed)
     np.random.seed(seed)