File size: 1,303 Bytes
6796365 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | import numpy as np
import os
import torch
from rdkit import Chem
from rdkit.Chem import AllChem, DataStructs
def smiles_to_ecfp(smiles, radius=2, n_bits=1024):
mol = Chem.MolFromSmiles(smiles)
if mol is None:
return np.zeros(n_bits)
fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius, nBits=n_bits)
arr = np.zeros(n_bits, dtype=int)
DataStructs.ConvertToNumpyArray(fp, arr)
return arr
class ModelWrapper:
def __init__(self, model_name: str = None):
self.model = None
model_name = os.path.join(os.environ.get("MODELS_DIR"), model_name)
print(model_name)
if model_name and os.path.exists(model_name):
try:
self.model = torch.load(model_name, map_location="cpu", weights_only=False)
except Exception as e:
print(e)
self.model = None
print(self.model)
self.featurizer = smiles_to_ecfp
def predict(self, X):
X = self.featurizer(X)
X = np.asarray(X, dtype=float)
# self.model.eval()
with torch.no_grad():
t = torch.tensor(X, dtype=torch.float32)
out = self.model(t)
# print(out.cpu().numpy().item())
return out.cpu().numpy().item()
|