aakothari/tox21_deepberta

1"""2Inference for the Tox21 leaderboard Space.3 4Contract (from the leaderboard paper, Appendix A.1 and Section 4):5  - predict() takes a list of SMILES strings6  - returns {smiles: {task_name: score}} for ALL 12 tasks7  - scores are floats in [0, 1]8  - EVERY input molecule must get a prediction for EVERY task. The leaderboard9    validates all molecule-target pairs are present and errors out otherwise.10    Never drop a molecule, no matter what the preprocessing does.11"""12 13import os14import torch15import torch.nn as nn16import deepsmiles17from rdkit import Chem, RDLogger18from transformers import AutoTokenizer, AutoModel19 20RDLogger.DisableLog("rdApp.*")21 22TASKS = [23    "NR-AR", "NR-AR-LBD", "NR-AhR", "NR-Aromatase", "NR-ER", "NR-ER-LBD",24    "NR-PPAR-gamma", "SR-ARE", "SR-ATAD5", "SR-HSE", "SR-MMP", "SR-p53",25]26 27CHECKPOINT = os.environ.get("CHECKPOINT_DIR", "checkpoints/deepberta-tox21")28CONVERTER = deepsmiles.Converter(rings=True, branches=True)29DEVICE = "cuda" if torch.cuda.is_available() else "cpu"30 31# Neutral score for molecules we cannot encode. 0.5 is uninformative for32# ROC-AUC rather than actively wrong. Any conversion failures MUST be33# reported in the model card -- silently scoring 0.5 on part of the test34# set is exactly the kind of thing reviewers need disclosed.35FALLBACK_SCORE = 0.536 37 38class DeepBERTaForTox21(nn.Module):39    def __init__(self, path, n_tasks=12, dropout=0.1):40        super().__init__()41        self.encoder = AutoModel.from_pretrained(path)42        self.dropout = nn.Dropout(dropout)43        self.classifier = nn.Linear(self.encoder.config.hidden_size, n_tasks)44 45    def forward(self, input_ids, attention_mask):46        out = self.encoder(input_ids=input_ids, attention_mask=attention_mask)47        mask = attention_mask.unsqueeze(-1).float()48        pooled = (out.last_hidden_state * mask).sum(1) / mask.sum(1).clamp(min=1e-9)49        return self.classifier(self.dropout(pooled))50 51 52_tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)53_model = DeepBERTaForTox21(CHECKPOINT).to(DEVICE)54_model.load_state_dict(55    torch.load(os.path.join(CHECKPOINT, "model.pt"), map_location=DEVICE,56               weights_only=True)57)58_model.eval()59 60 61def _encode(smiles):62    """SMILES -> DeepSMILES. Returns None if unconvertible."""63    try:64        mol = Chem.MolFromSmiles(smiles)65        if mol is None:66            return None67        return CONVERTER.encode(Chem.MolToSmiles(mol))68    except Exception:69        return None70 71 72def predict(smiles_list, batch_size=32):73    results = {s: {t: FALLBACK_SCORE for t in TASKS} for s in smiles_list}74 75    encodable = [(s, _encode(s)) for s in smiles_list]76    usable = [(s, d) for s, d in encodable if d is not None]77    n_failed = len(smiles_list) - len(usable)78    if n_failed:79        print(f"WARNING: {n_failed}/{len(smiles_list)} molecules failed "80              f"DeepSMILES conversion; returning {FALLBACK_SCORE} for these.")81 82    with torch.no_grad():83        for i in range(0, len(usable), batch_size):84            chunk = usable[i:i + batch_size]85            enc = _tokenizer(86                [d for _, d in chunk], padding=True, truncation=True,87                max_length=256, return_tensors="pt",88            ).to(DEVICE)89            # sigmoid: the leaderboard requires values in [0, 1], not logits90            probs = torch.sigmoid(_model(**enc)).cpu().numpy()91            for (smi, _), row in zip(chunk, probs):92                results[smi] = {t: float(p) for t, p in zip(TASKS, row)}93 94    # The leaderboard requires a prediction for every molecule-target pair it95    # sent. The test set is 647 molecules but only 645 unique structures, so96    # check against the unique inputs -- and verify none went missing.97    missing = set(smiles_list) - set(results)98    assert not missing, f"no prediction for {len(missing)} molecules: {list(missing)[:5]}"99    for s, d in results.items():100        assert set(d) == set(TASKS), f"incomplete predictions for {s}"101        assert all(0.0 <= v <= 1.0 for v in d.values()), f"out of range for {s}"102 103    print(f"returned predictions for {len(results)} unique structures "104          f"from {len(smiles_list)} inputs, {len(TASKS)} tasks each")105 106    return results107 108 109if __name__ == "__main__":110    # awkward cases worth checking before deploying: salts, stereochemistry,111    # charged species, very small molecules, and deliberate garbage112    tests = [113        "CCO",114        "CC(=O)Oc1ccccc1C(=O)O",115        "C[C@H](N)C(=O)O",116        "[Na+].[Cl-]",117        "c1ccc2c(c1)ccc1ccccc12",118        "not_a_molecule",119    ]120    for smi, preds in predict(tests).items():121        print(f"{smi:35s} {preds['NR-AR']:.3f}  {preds['SR-MMP']:.3f}")