| 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) |
|
|
| |
|
|
| |
| with torch.no_grad(): |
| t = torch.tensor(X, dtype=torch.float32) |
| out = self.model(t) |
| |
| return out.cpu().numpy().item() |
|
|
