app.py

from bmfm_sm.api.smmv_api import SmallMoleculeMultiViewModel
from bmfm_sm.core.data_modules.namespace import LateFusionStrategy
from bmfm_sm.api.dataset_registry import DatasetRegistry

import gradio as gr


examples = [
    ["CC(C)CC1=CC=C(C=C1)C(C)C(=O)O", "BACE"],
    ["CC(C)CC1=CC=C(C=C1)C(C)C(=O)O", "BBBP"],
    ["[N+](=O)([O-])[O-]", "CLINTOX"],
    ["OCC3OC(OCC2OC(OC(C#N)c1ccccc1)C(O)C(O)C2O)C(O)C(O)C3O", "ESOL"],
    ["CN(C)C(=O)c1ccc(cc1)OC", "FREESOLV"],
    ["CC(C)CC1=CC=C(C=C1)C(C)C(=O)O", "HIV"],
    ["Cn1c(CN2CCN(CC2)c3ccc(Cl)cc3)nc4ccccc14", "LIPOPHILICITY"],
    ["Cc1cccc(N2CCN(C(=O)C34CC5CC(CC(C5)C3)C4)CC2)c1C", "MUV"],
    ["C([H])([H])([H])[H]", "QM7"],
    ["C(CNCCNCCNCCN)N", "SIDER"],
    ["CCOc1ccc2nc(S(N)(=O)=O)sc2c1", "TOX21"],
    ["CSc1nc(N)nc(-c2cccc(-c3ccc4[nH]ccc4c3)c2)n1", "Pretrained"]
]

base_huggingface_path = 'ibm/biomed.sm.mv-te-84m'
finetuned_huggingface_path = "-MoleculeNet-ligand_scaffold-"

available_datasets = {
    "BACE": "ibm/biomed.sm.mv-te-84m-MoleculeNet-ligand_scaffold-BACE-101",
    "BBBP": "ibm/biomed.sm.mv-te-84m-MoleculeNet-ligand_scaffold-BBBP-101",
    "CLINTOX": "ibm/biomed.sm.mv-te-84m-MoleculeNet-ligand_scaffold-CLINTOX-101",
    "ESOL": "ibm/biomed.sm.mv-te-84m-MoleculeNet-ligand_scaffold-ESOL-101",
    "FREESOLV": "ibm/biomed.sm.mv-te-84m-MoleculeNet-ligand_scaffold-FREESOLV-101",
    "HIV": "ibm/biomed.sm.mv-te-84m-MoleculeNet-ligand_scaffold-HIV-101",
    "LIPOPHILICITY": "ibm/biomed.sm.mv-te-84m-MoleculeNet-ligand_scaffold-LIPOPHILICITY-101",
    "MUV": "ibm/biomed.sm.mv-te-84m-MoleculeNet-ligand_scaffold-MUV-101",
    "QM7": "ibm/biomed.sm.mv-te-84m-MoleculeNet-ligand_scaffold-QM7-101",
    "SIDER": "ibm/biomed.sm.mv-te-84m-MoleculeNet-ligand_scaffold-SIDER-101",
    "TOX21": "ibm/biomed.sm.mv-te-84m-MoleculeNet-ligand_scaffold-TOX21-101",
}


class PretrainedSMMVPipeline:
    def __init__(self, pretrained_model_name_or_path: str):
        self.model = SmallMoleculeMultiViewModel.from_pretrained(
            LateFusionStrategy.ATTENTIONAL,
            model_path=pretrained_model_name_or_path,
            huggingface=True
        )

    def __call__(self, smiles: str) -> float:
        emb = SmallMoleculeMultiViewModel.get_embeddings(
            smiles=smiles,
            pretrained_model=self.model
        )
        return str(emb.tolist())


class FinetunedSMMVPipeline:
    def __init__(self, dataset:str, pretrained_model_name_or_path: str):
        dataset_registry = DatasetRegistry()
        self.ds = dataset_registry.get_dataset_info(dataset)
        self.model = SmallMoleculeMultiViewModel.from_finetuned(
            self.ds,
            model_path=pretrained_model_name_or_path,
            inference_mode=True,
            huggingface=True
        )

    def __call__(self, smiles: str) -> float:
        prediction = SmallMoleculeMultiViewModel.get_predictions(
            smiles,
            self.ds,
            finetuned_model=self.model
        )
        return str(prediction.tolist())


def deploy():
    print(f"Loading checkpoint: Pretrained from {base_huggingface_path}")
    pipeline_pretrained = PretrainedSMMVPipeline(base_huggingface_path)

    pipelines_finetuned = {}
    pipelines_finetuned["Pretrained"] = pipeline_pretrained

    for dataset, huggingface_path in available_datasets.items():
        print(f"Loading checkpoint: {dataset} from {huggingface_path}")
        pipelines_finetuned[dataset] = FinetunedSMMVPipeline(
            dataset=dataset,
            pretrained_model_name_or_path=huggingface_path
        )
    
    def pipeline(
        smiles: str,
        dataset: str
    ):
        return pipelines_finetuned[dataset](smiles)
    
    smiles_input = gr.Textbox(placeholder="SMILES", label="SMILES")
    datasets_input = gr.Dropdown(
        choices=list(pipelines_finetuned.keys()),
        label="Checkpoint",
    )
    text_output = gr.Textbox(
        max_lines=10,
        label="Prediction",
    )

    gradio_app = gr.Interface(
        pipeline,
        inputs=[smiles_input, datasets_input],
        outputs=text_output,
        examples=examples,
        examples_per_page=20,
        title="ibm/biomed.sm.mv-te-84m property prediction tasks",
        description="Predictions for Pretrained show embedding vector of base model. Predictions for datasets show output of model finetuned on that task",
    )

    gradio_app.launch()

if __name__ == "__main__":
    deploy()