Spaces:
Running
Running
| import logging | |
| import pathlib | |
| import gradio as gr | |
| import numpy as np | |
| import pandas as pd | |
| from gt4sd.properties.molecules import MOLECULE_PROPERTY_PREDICTOR_FACTORY | |
| from utils import draw_grid_predict | |
| logger = logging.getLogger(__name__) | |
| logger.addHandler(logging.NullHandler()) | |
| MOLFORMER_VERSIONS = { | |
| "molformer_classification": ["bace", "bbbp", "hiv"], | |
| "molformer_regression": [ | |
| "alpha", | |
| "cv", | |
| "g298", | |
| "gap", | |
| "h298", | |
| "homo", | |
| "lipo", | |
| "lumo", | |
| "mu", | |
| "r2", | |
| "u0", | |
| ], | |
| "molformer_multitask_classification": ["clintox", "sider", "tox21"], | |
| } | |
| REMOVE = ["docking", "docking_tdc", "molecule_one", "askcos", "plogp"] | |
| REMOVE.extend(["similarity_seed", "activity_against_target", "organtox"]) | |
| REMOVE.extend(MOLFORMER_VERSIONS.keys()) | |
| MODEL_PROP_DESCRIPTION = { | |
| "Tox21": "NR-AR, NR-AR-LBD, NR-AhR, NR-Aromatase, NR-ER, NR-ER-LBD, NR-PPAR-gamma, SR-ARE, SR-ATAD5, SR-HSE, SR-MMP, SR-p53", | |
| "Sider": "Hepatobiliary disorders,Metabolism and nutrition disorders,Product issues,Eye disorders,Investigations,Musculoskeletal disorders,Gastrointestinal disorders,Social circumstances,Immune system disorders,Reproductive system and breast disorders,Bening & malignant,General disorders,Endocrine disorders,Surgical & medical procedures,Vascular disorders,Blood & lymphatic disorders,Skin & subcutaneous disorders,Congenital & genetic disorders,Infections,Respiratory & thoracic disorders,Psychiatric disorders,Renal & urinary disorders,Pregnancy conditions,Ear disorders,Cardiac disorders,Nervous system disorders,Injury & procedural complications", | |
| "Clintox": "FDA approval, Clinical trial failure", | |
| } | |
| def main(property: str, smiles: str, smiles_file: str): | |
| if "Molformer" in property: | |
| version = property.split(" ")[-1].split("(")[-1].split(")")[0] | |
| property = property.split(" ")[0] | |
| algo, config = MOLECULE_PROPERTY_PREDICTOR_FACTORY[property.lower()] | |
| kwargs = ( | |
| {"algorithm_version": "v0"} if property in MODEL_PROP_DESCRIPTION.keys() else {} | |
| ) | |
| if property.lower() in MOLFORMER_VERSIONS.keys(): | |
| kwargs["algorithm_version"] = version | |
| model = algo(config(**kwargs)) | |
| if smiles != "" and smiles_file is not None: | |
| raise ValueError("Pass either smiles or smiles_file, not both.") | |
| elif smiles != "": | |
| smiles = [smiles] | |
| elif smiles_file is not None: | |
| smiles = pd.read_csv(smiles_file.name, header=None, sep="\t")[0].tolist() | |
| props = np.array(list(map(model, smiles))).round(2) | |
| # Expand to 2D array if needed | |
| if len(props.shape) == 1: | |
| props = np.expand_dims(np.array(props), -1) | |
| if property in MODEL_PROP_DESCRIPTION.keys(): | |
| property_names = MODEL_PROP_DESCRIPTION[property].split(",") | |
| else: | |
| property_names = [property] | |
| return draw_grid_predict( | |
| smiles, props, property_names=property_names, domain="Molecules" | |
| ) | |
| if __name__ == "__main__": | |
| # Preparation (retrieve all available algorithms) | |
| properties = list(MOLECULE_PROPERTY_PREDICTOR_FACTORY.keys())[::-1] | |
| for prop in REMOVE: | |
| prop_to_idx = dict(zip(properties, range(len(properties)))) | |
| properties.pop(prop_to_idx[prop]) | |
| properties = list(map(lambda x: x.capitalize(), properties)) | |
| # MolFormer options | |
| for key in MOLFORMER_VERSIONS.keys(): | |
| properties.extend( | |
| [f"{key.capitalize()} ({version})" for version in MOLFORMER_VERSIONS[key]] | |
| ) | |
| # Load metadata | |
| metadata_root = pathlib.Path(__file__).parent.joinpath("model_cards") | |
| examples = [ | |
| ["Qed", "", metadata_root.joinpath("examples.smi")], | |
| [ | |
| "Esol", | |
| "CN1CCN(CCCOc2ccc(N3C(=O)C(=Cc4ccc(Oc5ccc([N+](=O)[O-])cc5)cc4)SC3=S)cc2)CC1", | |
| None, | |
| ], | |
| ] | |
| with open(metadata_root.joinpath("article.md"), "r") as f: | |
| article = f.read() | |
| with open(metadata_root.joinpath("description.md"), "r") as f: | |
| description = f.read() | |
| demo = gr.Interface( | |
| fn=main, | |
| title="Molecular properties", | |
| inputs=[ | |
| gr.Dropdown(properties, label="Property", value="Scscore"), | |
| gr.Textbox( | |
| label="Single SMILES", | |
| placeholder="CC(C#C)N(C)C(=O)NC1=CC=C(Cl)C=C1", | |
| lines=1, | |
| ), | |
| gr.File( | |
| file_types=[".smi"], | |
| label="Multiple SMILES (tab-separated, `.smi` file)", | |
| ), | |
| ], | |
| outputs=gr.HTML(label="Output"), | |
| article=article, | |
| description=description, | |
| examples=examples, | |
| ) | |
| demo.launch(debug=True, show_error=True, share=True) | |