"""Forward EI predictors. All public prediction outputs are RAW relative intensity.""" from __future__ import annotations import warnings from pathlib import Path from typing import Optional import numpy as np import torch from torch import nn from torch.nn import functional as F from rdkit import Chem, DataStructs from rdkit.Chem import Descriptors, rdFingerprintGenerator def choose_device(value="auto"): if value == "auto": value = "cuda" if torch.cuda.is_available() else "cpu" device = torch.device(value) if device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError("CUDA requested but unavailable. Check the AutoDL PyTorch image / GPU.") return device def molecular_features(smiles, n_bits=2048, radius=2): generator = rdFingerprintGenerator.GetMorganGenerator(radius=radius, fpSize=n_bits) count_rows, bit_rows, mass_rows = [], [], [] for value in smiles: molecule = Chem.MolFromSmiles(str(value)) if molecule is None: raise ValueError(f"Invalid SMILES: {value!r}") counts = np.zeros(n_bits, dtype=np.float32) bits = np.zeros(n_bits, dtype=np.float32) DataStructs.ConvertToNumpyArray(generator.GetCountFingerprint(molecule), counts) DataStructs.ConvertToNumpyArray(generator.GetFingerprint(molecule), bits) count_rows.append(counts) bit_rows.append(bits) mass_rows.append(Descriptors.ExactMolWt(molecule)) if not count_rows: return np.empty((0, n_bits), np.float32), np.empty((0, n_bits), np.float32), np.empty(0, np.float32) return np.stack(count_rows), np.stack(bit_rows), np.asarray(mass_rows, dtype=np.float32) class ResidualBlock(nn.Module): def __init__(self, width, dropout): super().__init__() self.block = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, width), nn.GELU(), nn.Dropout(dropout), nn.Linear(width, width)) def forward(self, inputs): return inputs + self.block(inputs) class BidirectionalSpectrumHead(nn.Module): """NEIMS-inspired mass/neutral-loss head; output represents sqrt(intensity).""" def __init__(self, input_dim, output_dim, hidden_dim=512, blocks=2, dropout=0.1, isotope_margin=4): super().__init__() self.output_dim = int(output_dim) self.isotope_margin = int(isotope_margin) self.body = nn.Sequential(nn.Linear(input_dim, hidden_dim), nn.GELU(), *[ResidualBlock(hidden_dim, dropout) for _ in range(blocks)], nn.LayerNorm(hidden_dim)) self.forward_head = nn.Linear(hidden_dim, output_dim) self.reverse_head = nn.Linear(hidden_dim, output_dim) self.gate = nn.Linear(hidden_dim, output_dim) self.register_buffer("bins", torch.arange(output_dim), persistent=False) def forward(self, embeddings, exact_mass): hidden = self.body(embeddings) forward = F.softplus(self.forward_head(hidden)) reverse = F.softplus(self.reverse_head(hidden)) center = torch.floor(exact_mass.float() + 0.5).long() reverse_index = center[:, None] + self.isotope_margin - self.bins[None, :] reverse_valid = (reverse_index >= 0) & (reverse_index < self.output_dim) aligned = reverse.gather(1, reverse_index.clamp(0, self.output_dim - 1)) * reverse_valid gate = torch.sigmoid(self.gate(hidden)) prediction = gate * forward + (1.0 - gate) * aligned valid = (self.bins[None, :] > 0) & (self.bins[None, :] <= center[:, None] + self.isotope_margin) return prediction * valid class SafeBatchNorm(nn.BatchNorm1d): def forward(self, values): if self.training and values.shape[0] == 1: return F.batch_norm(values, self.running_mean, self.running_var, self.weight, self.bias, training=False, momentum=self.momentum, eps=self.eps) return super().forward(values) class NEIMSResidualBlock(nn.Module): def __init__(self, width, bottleneck_factor, dropout): super().__init__() inner = max(1, int(width * bottleneck_factor)) self.block = nn.Sequential(SafeBatchNorm(width), nn.ReLU(), nn.Dropout(dropout), nn.Linear(width, inner), SafeBatchNorm(inner), nn.ReLU(), nn.Linear(inner, width)) def forward(self, values): return values + self.block(values) class NEIMSReimplementation(nn.Module): """Independent implementation aligned to the published NEIMS-style MLP defaults. 4096 raw Morgan count features, width 2000, seven BN/ReLU residual bottlenecks, dropout .25, gated forward/reverse outputs. Shared split, sqrt-cosine loss and 0..750 bins differ from the original TensorFlow/NIST training protocol. """ def __init__(self, config): super().__init__() self.config = config width = config.get("hidden_dim", 2000) dimension = config["spectrum_dim"] self.isotope_margin = config.get("isotope_margin", 5) self.input = nn.Sequential(SafeBatchNorm(config.get("fingerprint_bits", 4096)), nn.Linear(config.get("fingerprint_bits", 4096), width), nn.ReLU()) self.blocks = nn.Sequential(*[NEIMSResidualBlock(width, config.get("bottleneck_factor", .5), config.get("dropout", .25)) for _ in range(config.get("residual_blocks", 7))], SafeBatchNorm(width), nn.ReLU()) self.forward_head = nn.Linear(width, dimension) self.reverse_head = nn.Linear(width, dimension) self.gate = nn.Linear(width, dimension) self.register_buffer("bins", torch.arange(dimension), persistent=False) def forward(self, inputs, exact_mass): features = self.blocks(self.input(inputs["fingerprints"])) direct = self.forward_head(features) reverse = self.reverse_head(features) ceiling = torch.floor(exact_mass.float() + .5).long() + self.isotope_margin source = ceiling[:, None] - self.bins[None, :] valid = (source >= 0) & (source < len(self.bins)) & (self.bins[None, :] > 0) aligned = reverse.gather(1, source.clamp(0, len(self.bins) - 1)) * valid gate = torch.sigmoid(self.gate(features)) direct_valid = (self.bins[None, :] > 0) & (self.bins[None, :] <= ceiling[:, None]) return F.relu(gate * direct + (1 - gate) * aligned) * direct_valid class MolformerSpectrumModel(nn.Module): def __init__(self, config, model_dir_override=None): super().__init__() from transformers import AutoModel, AutoTokenizer from peft import LoraConfig, get_peft_model self.config = config model_dir = Path(model_dir_override or config.get("model_dir", "models/molformer")) if not model_dir.is_dir() or not (model_dir / "config.json").is_file(): raise FileNotFoundError(f"MoLFormer files missing in {model_dir}. Run download_resources.py --molformer first.") self.tokenizer = AutoTokenizer.from_pretrained(str(model_dir), trust_remote_code=True, local_files_only=True) backbone = AutoModel.from_pretrained(str(model_dir), trust_remote_code=True, local_files_only=True, deterministic_eval=True) expected = config.get("lora_targets", ["query", "value"]) observed = {name.rsplit(".", 1)[-1] for name, module in backbone.named_modules() if isinstance(module, nn.Linear)} if not set(expected).issubset(observed): raise RuntimeError(f"LoRA modules {expected} not found. Available linear leaf names: {sorted(observed)}") lora = LoraConfig(r=config.get("lora_r", 8), lora_alpha=config.get("lora_alpha", 16), lora_dropout=config.get("lora_dropout", 0.05), target_modules=expected, bias="none") self.encoder = get_peft_model(backbone, lora) self.head = BidirectionalSpectrumHead( backbone.config.hidden_size, config["spectrum_dim"], config.get("hidden_dim", 512), config.get("residual_blocks", 2), config.get("dropout", 0.1), config.get("isotope_margin", 4)) def tokenize(self, smiles): # Check rather than silently truncate molecules. encoded = self.tokenizer(list(smiles), padding=True, truncation=False, return_tensors="pt") limit = self.config.get("max_tokens", 256) if encoded["input_ids"].shape[1] > limit: raise ValueError(f"SMILES exceeds max_tokens={limit}; review long structures or raise max_tokens explicitly.") return {key: value for key, value in encoded.items() if key in {"input_ids", "attention_mask"}} def forward(self, inputs, exact_mass): encoded = self.encoder(**inputs, return_dict=True) mask = inputs["attention_mask"].unsqueeze(-1).to(encoded.last_hidden_state.dtype) pooled = (encoded.last_hidden_state * mask).sum(1) / mask.sum(1).clamp_min(1) return self.head(pooled, exact_mass) def build_model(config, model_dir_override=None): if config["model_type"] == "neims_reimpl": return NEIMSReimplementation(config) if config["model_type"] == "molformer": return MolformerSpectrumModel(config, model_dir_override) raise ValueError(f"Unknown model_type: {config['model_type']}") def load_checkpoint(path, map_location="cpu"): # Only load checkpoints produced by this package or another trusted source. return torch.load(path, map_location=map_location, weights_only=False) def raw_intensity(sqrt_prediction): raw = sqrt_prediction.float().clamp_min(0).square() return raw / raw.amax(dim=-1, keepdim=True).clamp_min(1e-12) class SpectrumPredictor: def __init__(self, checkpoint, device="auto", model_dir_override=None): self.device = choose_device(device) saved = load_checkpoint(checkpoint) self.config = saved["config"] self.model = build_model(self.config, model_dir_override) self.model.load_state_dict(saved["model_state"], strict=True) self.model.to(self.device).eval() if "train_keys" not in saved: raise ValueError("Checkpoint lacks training identity audit (train_keys)") self.train_keys = set(saved["train_keys"]) self.metadata = {key: value for key, value in saved.items() if key not in {"model_state", "optimizer_state", "rng_state", "train_keys"}} @torch.inference_mode() def __call__(self, smiles, exact_mass=None, batch_size=32): smiles = list(smiles) if batch_size < 1: raise ValueError("batch_size must be positive") if not smiles: return np.empty((0, self.config["spectrum_dim"]), dtype=np.float32) supplied_mass = None if exact_mass is None else np.asarray(exact_mass, dtype=np.float32) if supplied_mass is not None and supplied_mass.shape != (len(smiles),): raise ValueError("exact_mass must have one value per SMILES") results = [] for start in range(0, len(smiles), batch_size): batch_smiles = smiles[start:start + batch_size] counts, _, calculated = molecular_features(batch_smiles, self.config.get("fingerprint_bits", 2048), self.config.get("fingerprint_radius", 2)) mass = calculated if supplied_mass is None else supplied_mass[start:start + batch_size] if not np.isfinite(mass).all() or (mass <= 0).any(): raise ValueError("Non-finite or nonpositive exact_mass") if self.config["model_type"] == "neims_reimpl": inputs = {"fingerprints": torch.from_numpy(counts)} else: inputs = self.model.tokenize(batch_smiles) inputs = {key: value.to(self.device) for key, value in inputs.items()} values = self.model(inputs, torch.as_tensor(mass, device=self.device)) results.append(raw_intensity(values).cpu().numpy()) return np.concatenate(results) def load_predictor(checkpoint, device="auto", model_dir_override=None): return SpectrumPredictor(checkpoint, device, model_dir_override)