Update tool/comget/generator.py

tool/comget/generator.py

@@ -5,13 +5,9 @@ import math
 from tqdm import tqdm
 import argparse
 from .model import GPT, GPTConfig
-import pandas as pd
 import torch
-import numpy as np
-
-from .moses.utils import get_mol
+import numpy as np 
 import re
-
 import json
 from rdkit.Chem import RDConfig
 from torch.nn import functional as F
@@ -23,6 +19,23 @@ from rdkit import Chem
 import os
 os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
 
+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 top_k_logits(logits, k):
     v, ix = torch.topk(logits, k)
     out = logits.clone()