import logging import math import os import pickle import sys import warnings from pathlib import Path import numpy as np import torch import torch.nn as nn from lightning import pytorch as pl from rdkit import Chem, RDLogger from transformers import AutoModelWithLMHead, AutoTokenizer, EsmModel REPO_ROOT = Path(__file__).resolve().parents[1] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) if str(Path(__file__).resolve().parent) not in sys.path: sys.path.insert(0, str(Path(__file__).resolve().parent)) from BindEvaluator_models import * # noqa: E402,F403 from DeepDTAGen_models import * # noqa: E402,F403 from paths import peptiverse_root # noqa: E402 os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error") os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1") torch.set_float32_matmul_precision("high") warnings.filterwarnings("ignore", category=FutureWarning) warnings.filterwarnings("ignore", category=UserWarning, message=r".*predict_dataloader.*many workers.*") warnings.filterwarnings("ignore", message=r"Dropping last batch of size .*") RDLogger.DisableLog("rdApp.*") logging.basicConfig(level=logging.ERROR, force=True) for name in ["lightning", "pytorch_lightning", "lightning.pytorch", "chemprop", "rdkit", "urllib3", "torch"]: logging.getLogger(name).setLevel(logging.ERROR) _AMIDE_SMARTS = Chem.MolFromSmarts("[CX3](=[OX1])[NX3]") _CARBONYL_C_SMARTS = Chem.MolFromSmarts("[CX3](=[OX1])") _DIPEPTIDE_SMARTS = Chem.MolFromSmarts("[CX3](=[OX1])N[#6X4][CX3](=[OX1])N") _PRED = None def get_peptiverse_predictor(): """Lazy PeptiVerse loader. Clone PeptiVerse and/or set PEPTIVERSE_ROOT first.""" global _PRED if _PRED is not None: return _PRED root = peptiverse_root() if str(root) not in sys.path: sys.path.insert(0, str(root)) from inference import PeptiVersePredictor device = "cuda" if torch.cuda.is_available() else "cpu" _PRED = PeptiVersePredictor( manifest_path=root / "best_models.txt", classifier_weight_root=root, device=device, ) return _PRED def _amide_bond_indices(mol, ignore_ring_amides=False): ids = set() for c_idx, _, n_idx in mol.GetSubstructMatches(_AMIDE_SMARTS): b = mol.GetBondBetweenAtoms(c_idx, n_idx) if b and b.GetBondType() == Chem.rdchem.BondType.SINGLE: if ignore_ring_amides and b.IsInRing(): continue ids.add(b.GetIdx()) return ids def _carbonyl_c_indices(mol): return {m[0] for m in mol.GetSubstructMatches(_CARBONYL_C_SMARTS)} def _carbonyl_neighbor_stats(mol, c_indices): stats = {"total": 0, "with_N": 0, "with_O": 0, "with_S": 0, "pure_amide": 0} for c_idx in c_indices: c = mol.GetAtomWithIdx(c_idx) stats["total"] += 1 hasN = hasO = hasS = False for b in c.GetBonds(): if b.GetBondType() != Chem.rdchem.BondType.SINGLE: continue z = b.GetOtherAtom(c).GetAtomicNum() if z == 7: hasN = True elif z == 8: hasO = True elif z == 16: hasS = True stats["with_N"] += int(hasN) stats["with_O"] += int(hasO) stats["with_S"] += int(hasS) if hasN and not (hasO or hasS): stats["pure_amide"] += 1 return stats def _adjacent_amide_pairs(mol): centers = set() for match in mol.GetSubstructMatches(_DIPEPTIDE_SMARTS): centers.add(match[3]) return len(centers) def analyze_peptide_likeness(smiles: str, ignore_ring_amides: bool = False, amide_density_target: float = 0.12): mol = Chem.MolFromSmiles(smiles) if mol is None: raise ValueError(f"Invalid SMILES: {smiles}") n_heavy_atoms = mol.GetNumAtoms() n_heavy_bonds = mol.GetNumBonds() amide_bonds = _amide_bond_indices(mol, ignore_ring_amides=ignore_ring_amides) n_amide_bonds = len(amide_bonds) carbonyl_cs = _carbonyl_c_indices(mol) cstats = _carbonyl_neighbor_stats(mol, carbonyl_cs) total_carb = max(1, cstats["total"]) f1 = cstats["with_N"] / total_carb amide_per_atom = n_amide_bonds / max(1, n_heavy_atoms) f2 = min(1.0, amide_per_atom / max(1e-8, amide_density_target)) f3 = 1.0 - math.exp(-_adjacent_amide_pairs(mol)) penalty = 1.0 - (cstats["pure_amide"] / total_carb) score = max(0.0, min(1.0, 0.55 * f1 + 0.25 * f2 + 0.20 * f3 - 0.25 * penalty)) return { "n_heavy_atoms": n_heavy_atoms, "n_heavy_bonds": n_heavy_bonds, "n_carbonyls": cstats["total"], "n_amide_bonds": n_amide_bonds, "amide_bond_ratio_all_bonds": n_amide_bonds / max(1, n_heavy_bonds), "acyl_N_fraction": f1, "pure_amide_fraction": cstats["pure_amide"] / total_carb, "amide_per_atom": amide_per_atom, "n_adjacent_amide_pairs": _adjacent_amide_pairs(mol), "peptide_likeness": score, } class Toxicity: """PeptiVerse SMILES toxicity. Returns non-toxicity so mix_score stays unchanged.""" def __call__(self, smiles_tokens, smiles_seqs): pred = get_peptiverse_predictor() scores = [] for seq in smiles_seqs: score = pred.predict_property("toxicity", col="smiles", input_str=seq)["score"] scores.append(1.0 - score) return "non_toxicity", scores class Solubility: def __call__(self, smiles_tokens, smiles_seqs): pred = get_peptiverse_predictor() scores = [pred.predict_property("solubility", col="smiles", input_str=seq)["score"] for seq in smiles_seqs] return "solubility", scores class Permeability: def __call__(self, smiles_tokens, smiles_seqs): pred = get_peptiverse_predictor() scores = [] for seq in smiles_seqs: score = pred.predict_property("permeability_pampa", col="smiles", input_str=seq)["score"] scores.append((score + 9.0) / (-4.0 + 9.0)) return "permeability", scores class Halflife: def __call__(self, smiles_tokens, smiles_seqs): pred = get_peptiverse_predictor() scores = [pred.predict_property("halflife", col="smiles", input_str=seq)["score"] for seq in smiles_seqs] return "halflife", scores class BindingAffinity: def __init__(self, prot_seq, device=None): self.prot_seq = prot_seq self.device = device def __call__(self, smiles_tokens, smiles_seqs): pred = get_peptiverse_predictor() scores = [] for seq in smiles_seqs: score = pred.predict_binding_affinity(col="smiles", target_seq=self.prot_seq, binder_str=seq)["affinity"] scores.append(min(10.0, score) / 10.0) return "affinity", scores class SmallMolecule: def __init__(self, protein_sequence, device): import chemprop self.chemprop = chemprop self.trainer = pl.Trainer(logger=False, enable_progress_bar=False, accelerator="cuda" if torch.cuda.is_available() else "cpu", devices=1) ckpt_dir = REPO_ROOT / "peptidomimetics" / "ckpt" / "admetica" self.models = self.load_models(ckpt_dir) model_path = REPO_ROOT / "peptidomimetics" / "ckpt" / "deepdtagen" / "deepdtagen_model_bindingdb.pth" tokenizer_path = REPO_ROOT / "peptidomimetics" / "ckpt" / "deepdtagen" / "bindingdb_tokenizer.pkl" with open(tokenizer_path, "rb") as f: tokenizer = pickle.load(f) self.deep_dta_gen = DeepDTAGen(tokenizer) self.deep_dta_gen.load_state_dict(torch.load(model_path, map_location=device, weights_only=False)) self.deep_dta_gen.to(device) self.deep_dta_gen.eval() self.protein_sequence = protein_sequence self.device = device def load_models(self, ckpt_dir): chemprop = self.chemprop toxicity_model = chemprop.models.MPNN.load_from_checkpoint(os.path.join(ckpt_dir, "ld50.ckpt")) solubility_model = chemprop.models.MPNN.load_from_checkpoint(os.path.join(ckpt_dir, "solubility.ckpt")) permeability_model = chemprop.models.MPNN.load_from_checkpoint(os.path.join(ckpt_dir, "caco2.ckpt")) halflife_model = chemprop.models.MPNN.load_from_checkpoint(os.path.join(ckpt_dir, "half-life.ckpt")) return toxicity_model, solubility_model, permeability_model, halflife_model def is_valid_smiles(self, smiles): try: return Chem.MolFromSmiles(smiles) is not None except Exception: return False def prediction(self, smiles_list, trainer, model): valid_smiles = [smi for smi in smiles_list if self.is_valid_smiles(smi)] valid_indices = [i for i, smi in enumerate(smiles_list) if self.is_valid_smiles(smi)] invalid_indices = [i for i in range(len(smiles_list)) if i not in valid_indices] if not valid_smiles: return np.full(len(smiles_list), "", dtype=object) test_data = [self.chemprop.data.MoleculeDatapoint.from_smi(smi) for smi in valid_smiles] featurizer = self.chemprop.featurizers.SimpleMoleculeMolGraphFeaturizer() test_dataset = self.chemprop.data.MoleculeDataset(test_data, featurizer=featurizer) test_loader = self.chemprop.data.build_dataloader(test_dataset, shuffle=False) with torch.no_grad(): predictions = trainer.predict(model, test_loader) predictions = [pred.item() for batch in predictions for pred in batch] for index in invalid_indices: predictions.insert(index, "") return predictions def non_toxicity_from_log10mgkg(self, scores, lo=1.0, hi=4.0): res = [] for score in scores: score = max(lo, min(hi, score)) res.append((score - lo) / (hi - lo)) return res def __call__(self, smiles_tokens, smiles_seqs): scores = [self.prediction(smiles_seqs, self.trainer, model) for model in self.models] non_toxicity = self.non_toxicity_from_log10mgkg(scores[0]) solubility = [(max(-12, min(2, score)) + 12) / 14 for score in scores[1]] permeability = [(max(-8, min(-3, score)) + 8) / 5 for score in scores[2]] halflife = scores[3] ratio = [analyze_peptide_likeness(smiles)["peptide_likeness"] for smiles in smiles_seqs] protein_sequences = [self.protein_sequence] * len(smiles_seqs) batch = process_latent_batch(smiles_seqs, protein_sequences).to(self.device) with torch.no_grad(): affinity = self.deep_dta_gen(batch).squeeze(-1) affinity = [min(10, score.item()) / 10 for score in affinity] return "small_molecule", { "non_toxicity": non_toxicity, "solubility": solubility, "permeability": permeability, "halflife": halflife, "affinity": affinity, "ratio": ratio, } def parse_motifs(motif: str) -> list: parts = motif.split(",") result = [] for part in parts: part = part.strip() if "-" in part: start, end = map(int, part.split("-")) result.extend(range(start, end + 1)) else: result.append(int(part)) print(f"Target Motifs: {result}") return torch.tensor(result) class BindEvaluator(pl.LightningModule): def __init__(self, cfg): super(BindEvaluator, self).__init__() self.esm_model = EsmModel.from_pretrained("facebook/esm2_t33_650M_UR50D").eval() for param in self.esm_model.parameters(): param.requires_grad = False self.chemberta_model = AutoModelWithLMHead.from_pretrained("seyonec/ChemBERTa_zinc250k_v2_40k").roberta.eval() for param in self.chemberta_model.parameters(): param.requires_grad = False self.repeated_module = RepeatedModule( cfg.model.n_layers, cfg.model.d_model, cfg.model.d_hidden, cfg.model.n_head, cfg.model.d_k, cfg.model.d_v, cfg.model.d_inner, dropout=cfg.model.dropout, ) self.final_attention_layer = MultiHeadAttentionSequence( cfg.model.n_head, cfg.model.d_model, cfg.model.d_k, cfg.model.d_v, dropout=cfg.model.dropout ) self.final_ffn = FFN(cfg.model.d_model, cfg.model.d_inner, dropout=cfg.model.dropout) self.output_projection_prot = nn.Linear(cfg.model.d_model, 1) def forward(self, binder_tokens, target_tokens): peptide_sequence = self.chemberta_model(**binder_tokens).last_hidden_state protein_sequence = self.esm_model(**target_tokens).last_hidden_state binder_mask = binder_tokens["attention_mask"] target_mask = target_tokens["attention_mask"] prot_enc, sequence_enc, sequence_attention_list, prot_attention_list, \ prot_seq_attention_list, seq_prot_attention_list = self.repeated_module( peptide_sequence, protein_sequence, peptide_mask=binder_mask, protein_mask=target_mask ) prot_enc, final_prot_seq_attention = self.final_attention_layer( prot_enc, sequence_enc, sequence_enc, key_padding_mask=binder_mask, query_padding_mask=target_mask, ) prot_enc = self.final_ffn(prot_enc, padding_mask=target_mask) prot_enc = self.output_projection_prot(prot_enc) return prot_enc class MotifModel: def __init__(self, cfg, target, motifs, device, specificity): self.cfg = cfg self.threshold = 0.918 self.device = device self.specificity = specificity self.chemberta_tokenizer = AutoTokenizer.from_pretrained("seyonec/ChemBERTa_zinc250k_v2_40k") self.esm_tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D") self.target = self.esm_tokenizer(target, return_tensors="pt").to(device) self.motifs = parse_motifs(motifs).to(device) ckpt = cfg.inference.ckpt if not os.path.isabs(ckpt): ckpt = str(REPO_ROOT / ckpt) self.bindevaluator = BindEvaluator.load_from_checkpoint(ckpt, cfg=cfg, map_location=device) def __call__(self, smiles_tokens, smiles_seqs): binder = self.chemberta_tokenizer( smiles_seqs, return_tensors="pt", padding=True, truncation=True, max_length=512 ).to(self.device) batch_size = binder["input_ids"].shape[0] L = self.target["input_ids"].shape[1] target = { "input_ids": self.target["input_ids"].expand(batch_size, L), "attention_mask": self.target["attention_mask"].expand(batch_size, L), } prediction = self.bindevaluator(binder, target).squeeze(-1) probs = torch.sigmoid(prediction) motif_scores = probs[:, self.motifs].mean(dim=1) if self.specificity: non_motif_probs = probs[:, [i for i in range(probs.shape[1]) if i not in self.motifs]] mask = non_motif_probs >= self.threshold count = mask.sum(dim=-1) specificity = 1 - count / (L - 2) return "motif", (motif_scores.tolist(), specificity.tolist()) return "motif", motif_scores.tolist()