Spaces:
Sleeping
Sleeping
| """ | |
| Inference for the Tox21 leaderboard Space. | |
| Contract (from the leaderboard paper, Appendix A.1 and Section 4): | |
| - predict() takes a list of SMILES strings | |
| - returns {smiles: {task_name: score}} for ALL 12 tasks | |
| - scores are floats in [0, 1] | |
| - EVERY input molecule must get a prediction for EVERY task. The leaderboard | |
| validates all molecule-target pairs are present and errors out otherwise. | |
| Never drop a molecule, no matter what the preprocessing does. | |
| """ | |
| import os | |
| import torch | |
| import torch.nn as nn | |
| import deepsmiles | |
| from rdkit import Chem, RDLogger | |
| from transformers import AutoTokenizer, AutoModel | |
| RDLogger.DisableLog("rdApp.*") | |
| TASKS = [ | |
| "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", | |
| ] | |
| CHECKPOINT = os.environ.get("CHECKPOINT_DIR", "checkpoints/deepberta-tox21") | |
| CONVERTER = deepsmiles.Converter(rings=True, branches=True) | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| # Neutral score for molecules we cannot encode. 0.5 is uninformative for | |
| # ROC-AUC rather than actively wrong. Any conversion failures MUST be | |
| # reported in the model card -- silently scoring 0.5 on part of the test | |
| # set is exactly the kind of thing reviewers need disclosed. | |
| FALLBACK_SCORE = 0.5 | |
| class DeepBERTaForTox21(nn.Module): | |
| def __init__(self, path, n_tasks=12, dropout=0.1): | |
| super().__init__() | |
| self.encoder = AutoModel.from_pretrained(path) | |
| self.dropout = nn.Dropout(dropout) | |
| self.classifier = nn.Linear(self.encoder.config.hidden_size, n_tasks) | |
| def forward(self, input_ids, attention_mask): | |
| out = self.encoder(input_ids=input_ids, attention_mask=attention_mask) | |
| mask = attention_mask.unsqueeze(-1).float() | |
| pooled = (out.last_hidden_state * mask).sum(1) / mask.sum(1).clamp(min=1e-9) | |
| return self.classifier(self.dropout(pooled)) | |
| _tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT) | |
| _model = DeepBERTaForTox21(CHECKPOINT).to(DEVICE) | |
| _model.load_state_dict( | |
| torch.load(os.path.join(CHECKPOINT, "model.pt"), map_location=DEVICE, | |
| weights_only=True) | |
| ) | |
| _model.eval() | |
| def _encode(smiles): | |
| """SMILES -> DeepSMILES. Returns None if unconvertible.""" | |
| try: | |
| mol = Chem.MolFromSmiles(smiles) | |
| if mol is None: | |
| return None | |
| return CONVERTER.encode(Chem.MolToSmiles(mol)) | |
| except Exception: | |
| return None | |
| def predict(smiles_list, batch_size=32): | |
| results = {s: {t: FALLBACK_SCORE for t in TASKS} for s in smiles_list} | |
| encodable = [(s, _encode(s)) for s in smiles_list] | |
| usable = [(s, d) for s, d in encodable if d is not None] | |
| n_failed = len(smiles_list) - len(usable) | |
| if n_failed: | |
| print(f"WARNING: {n_failed}/{len(smiles_list)} molecules failed " | |
| f"DeepSMILES conversion; returning {FALLBACK_SCORE} for these.") | |
| with torch.no_grad(): | |
| for i in range(0, len(usable), batch_size): | |
| chunk = usable[i:i + batch_size] | |
| enc = _tokenizer( | |
| [d for _, d in chunk], padding=True, truncation=True, | |
| max_length=256, return_tensors="pt", | |
| ).to(DEVICE) | |
| # sigmoid: the leaderboard requires values in [0, 1], not logits | |
| probs = torch.sigmoid(_model(**enc)).cpu().numpy() | |
| for (smi, _), row in zip(chunk, probs): | |
| results[smi] = {t: float(p) for t, p in zip(TASKS, row)} | |
| # The leaderboard requires a prediction for every molecule-target pair it | |
| # sent. The test set is 647 molecules but only 645 unique structures, so | |
| # check against the unique inputs -- and verify none went missing. | |
| missing = set(smiles_list) - set(results) | |
| assert not missing, f"no prediction for {len(missing)} molecules: {list(missing)[:5]}" | |
| for s, d in results.items(): | |
| assert set(d) == set(TASKS), f"incomplete predictions for {s}" | |
| assert all(0.0 <= v <= 1.0 for v in d.values()), f"out of range for {s}" | |
| print(f"returned predictions for {len(results)} unique structures " | |
| f"from {len(smiles_list)} inputs, {len(TASKS)} tasks each") | |
| return results | |
| if __name__ == "__main__": | |
| # awkward cases worth checking before deploying: salts, stereochemistry, | |
| # charged species, very small molecules, and deliberate garbage | |
| tests = [ | |
| "CCO", | |
| "CC(=O)Oc1ccccc1C(=O)O", | |
| "C[C@H](N)C(=O)O", | |
| "[Na+].[Cl-]", | |
| "c1ccc2c(c1)ccc1ccccc12", | |
| "not_a_molecule", | |
| ] | |
| for smi, preds in predict(tests).items(): | |
| print(f"{smi:35s} {preds['NR-AR']:.3f} {preds['SR-MMP']:.3f}") |