| import argparse |
| import hashlib |
| import importlib.util |
| import json |
| import os |
| import threading |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.distributed.tensor |
| from rdkit import Chem, DataStructs |
| from rdkit.Chem import Descriptors, rdFingerprintGenerator, rdMolDescriptors |
|
|
| from smart_core import AnalogRouter, predict_router |
|
|
|
|
| BASE_SHA256 = "2b5d2907b415aefd2f76bc0966bb1124a27c8e1c63fba60dafc54a9f8cc9f33e" |
| SOURCE_SHA256 = "3bcc82f45f761b453e75d716f4540c9a24a22c87637645373db2ea437d4bc394" |
| ROUTER_SHA256 = "4bf2b97bd618d4fd04b295a31db3f0c57b3d0be6684b02027dffdfac08966b32" |
| PROTOCOL = {"alpha": 0.25, "neighbors": 3, "minimum_similarity": 0.35, "similarity_power": 4} |
| WARNING_MESSAGES = { |
| "stereochemistry_ignored": "Stereochemistry is removed, consistently with the original connectivity-based model.", |
| "no_eligible_reference": "No eligible training analogue was found; SMART uses its structure-based prediction for this molecule.", |
| "reference_transfer_unavailable": "Selected reference peaks fall outside the allowed target-mass support; SMART uses its structure-based prediction.", |
| "seen_in_training": "This molecular connectivity occurs in the supervised training set. Its own measured spectrum is excluded from reference transfer.", |
| "limited_mass_range": "The molecular mass exceeds the output range. Only m/z 0–750 are predicted.", |
| } |
|
|
|
|
| def file_sha256(path): |
| digest = hashlib.sha256() |
| with Path(path).open("rb") as handle: |
| for block in iter(lambda: handle.read(1024 * 1024), b""): |
| digest.update(block) |
| return digest.hexdigest() |
|
|
|
|
| def molecular_identity(text): |
| if not isinstance(text, str) or not text.strip(): |
| raise ValueError("Enter one SMILES string.") |
| value = text.strip() |
| if len(value) > 4096 or len(value.split()) != 1: |
| raise ValueError("Enter one SMILES string without names, separators, or multiple lines.") |
| molecule = Chem.MolFromSmiles(value) |
| if molecule is None or not molecule.GetNumAtoms(): |
| raise ValueError("RDKit could not parse the SMILES string.") |
| if len(Chem.GetMolFrags(molecule)) != 1: |
| raise ValueError("Multicomponent structures are unsupported. Supply a single neutral molecule.") |
| if Chem.GetFormalCharge(molecule) != 0: |
| raise ValueError("Charged structures are unsupported. Supply the neutral molecular structure.") |
| if any(atom.GetIsotope() for atom in molecule.GetAtoms()): |
| raise ValueError("Isotope-labeled structures are unsupported.") |
| if any(atom.GetAtomicNum() == 0 for atom in molecule.GetAtoms()): |
| raise ValueError("Wildcard atoms are unsupported.") |
| if any(atom.GetAtomMapNum() for atom in molecule.GetAtoms()): |
| raise ValueError("Atom-mapped structures are unsupported.") |
| stereo = any(atom.GetChiralTag() != Chem.ChiralType.CHI_UNSPECIFIED for atom in molecule.GetAtoms()) |
| stereo = stereo or any(bond.GetStereo() != Chem.BondStereo.STEREONONE for bond in molecule.GetBonds()) |
| molecule = Chem.RemoveHs(molecule) |
| Chem.RemoveStereochemistry(molecule) |
| canonical = Chem.MolToSmiles(molecule, canonical=True, isomericSmiles=False) |
| molecule = Chem.MolFromSmiles(canonical) |
| inchikey = Chem.MolToInchiKey(molecule) |
| if not inchikey or len(inchikey.split("-")[0]) != 14: |
| raise ValueError("A molecular identity could not be generated from this structure.") |
| mass = float(Descriptors.ExactMolWt(molecule)) |
| if not np.isfinite(mass) or mass <= 0: |
| raise ValueError("The molecular exact mass is invalid.") |
| return { |
| "molecule": molecule, |
| "smiles": canonical, |
| "group_key": inchikey.split("-")[0], |
| "formula": rdMolDescriptors.CalcMolFormula(molecule), |
| "exact_mass": mass, |
| "stereochemistry_ignored": bool(stereo), |
| } |
|
|
|
|
| def l1_normalize(values): |
| values = np.asarray(values, dtype=np.float32) |
| if not np.isfinite(values).all() or (values < 0).any(): |
| raise ValueError("The predictor returned nonfinite or negative intensities.") |
| totals = values.sum(axis=-1, keepdims=True) |
| return np.divide(values, totals, out=np.zeros_like(values), where=totals > 0) |
|
|
|
|
| class SMARTPredictor: |
| def __init__(self, assets_dir="assets"): |
| self.assets_dir = Path(assets_dir).expanduser().resolve() |
| self._lock = threading.RLock() |
| self.manifest = self._verify_assets() |
| self.generator = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=2048, includeChirality=False) |
| self._load_bank() |
| os.environ["HF_HUB_OFFLINE"] = "1" |
| os.environ["TRANSFORMERS_OFFLINE"] = "1" |
| module_path = self.assets_dir / "runtime/src/models.py" |
| spec = importlib.util.spec_from_file_location("smart_frozen_v1_models", module_path) |
| module = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(module) |
| self.base = module.load_predictor( |
| str(self.assets_dir / "base.pt"), |
| device="cpu", |
| model_dir_override=str(self.assets_dir / "molformer"), |
| ) |
| if self.base.config.get("model_type") != "molformer" or self.base.config.get("spectrum_dim") != 751: |
| raise ValueError("The base checkpoint is not the frozen 751-bin MoLFormer EI predictor.") |
| if set(self.base.train_keys) != set(self.group_keys.tolist()): |
| raise ValueError("The reference bank does not match the base checkpoint's supervised training identities.") |
| if self.base.model.training: |
| raise ValueError("The base predictor is unexpectedly in training mode.") |
| self.router = AnalogRouter(751, fp_dim=2048) |
| state = torch.load(self.assets_dir / "router.pt", map_location="cpu", weights_only=True) |
| self.router.load_state_dict(state, strict=True) |
| self.router.eval() |
| if sum(parameter.numel() for parameter in self.router.parameters()) != 146019: |
| raise ValueError("The SMART transfer module has an unexpected parameter count.") |
| self.device = torch.device("cpu") |
| self.model_version = "SMART-frozen-seed42" |
|
|
| def _verify_assets(self): |
| path = self.assets_dir / "manifest.json" |
| if not path.is_file(): |
| raise FileNotFoundError("SMART assets are missing. Run the AutoDL export script and upload the complete assets folder.") |
| manifest = json.loads(path.read_text(encoding="utf-8")) |
| expected = {"schema_version": 1, "model": "SMART", "spectrum_dim": 751, "reference_count": 7382, "seed": 42} |
| for key, value in expected.items(): |
| if manifest.get(key) != value: |
| raise ValueError("Unsupported or incomplete SMART asset manifest: " + key) |
| if manifest.get("protocol") != PROTOCOL: |
| raise ValueError("The asset protocol differs from the frozen SMART method.") |
| files = manifest.get("files") |
| if not isinstance(files, dict) or not files: |
| raise ValueError("The asset manifest must contain SHA-256 hashes for all model files.") |
| pinned = {"base.pt": BASE_SHA256, "router.pt": ROUTER_SHA256, "runtime/src/models.py": SOURCE_SHA256} |
| for name, expected_hash in pinned.items(): |
| if files.get(name) != expected_hash: |
| raise ValueError("The manifest does not identify the original frozen asset: " + name) |
| required = set(pinned) | {"reference.npz", "molformer/config.json", "molformer/tokenizer_config.json"} |
| if not required.issubset(files): |
| raise ValueError("Required model assets are absent from the manifest.") |
| weights = [name for name in files if name.startswith("molformer/") and Path(name).suffix in {".bin", ".safetensors"}] |
| if not weights: |
| raise ValueError("The manifest contains no MoLFormer backbone weights.") |
| for name, expected_hash in files.items(): |
| relative = Path(name) |
| resolved = (self.assets_dir / relative).resolve() |
| if relative.is_absolute() or ".." in relative.parts or self.assets_dir not in resolved.parents: |
| raise ValueError("Unsafe relative asset path: " + name) |
| if not isinstance(expected_hash, str) or len(expected_hash) != 64: |
| raise ValueError("Invalid asset SHA-256: " + name) |
| if not resolved.is_file() or file_sha256(resolved) != expected_hash: |
| raise ValueError("Missing or modified SMART asset: " + name) |
| for folder in ("molformer", "runtime"): |
| for item in (self.assets_dir / folder).rglob("*"): |
| if item.is_file() and "__pycache__" not in item.parts: |
| name = item.relative_to(self.assets_dir).as_posix() |
| if name not in files: |
| raise ValueError("Unverified runtime or backbone file: " + name) |
| return manifest |
|
|
| def _load_bank(self): |
| with np.load(self.assets_dir / "reference.npz", allow_pickle=False) as archive: |
| required = {"group_key", "smiles", "features", "masses", "spectra"} |
| if not required.issubset(archive.files): |
| raise ValueError("The reference bank is missing required arrays.") |
| self.group_keys = np.asarray(archive["group_key"], dtype=str) |
| self.smiles = np.asarray(archive["smiles"], dtype=str) |
| self.features = np.asarray(archive["features"], dtype=np.float32) |
| self.masses = np.asarray(archive["masses"], dtype=np.float32) |
| self.spectra = np.asarray(archive["spectra"], dtype=np.float32) |
| count = self.manifest["reference_count"] |
| expected = ((self.group_keys, (count,)), (self.smiles, (count,)), (self.features, (count, 2048)), |
| (self.masses, (count,)), (self.spectra, (count, 751))) |
| if any(array.shape != shape for array, shape in expected): |
| raise ValueError("The reference bank has inconsistent array dimensions.") |
| if len(set(self.group_keys.tolist())) != count: |
| raise ValueError("The reference bank contains duplicate identities.") |
| if any(not np.isfinite(array).all() for array in (self.features, self.masses, self.spectra)): |
| raise ValueError("The reference bank contains nonfinite values.") |
| if not np.isin(self.features, [0, 1]).all() or (self.masses <= 0).any() or (self.spectra < 0).any(): |
| raise ValueError("Invalid reference fingerprints, masses, or intensities.") |
| if not np.allclose(self.spectra.sum(axis=1, dtype=np.float64), 1, atol=2e-6, rtol=0): |
| raise ValueError("Reference spectra must be nonempty L1-normalized consensus spectra.") |
| self.fingerprints = [] |
| self.connectivity = [] |
| for index, value in enumerate(self.smiles): |
| identity = molecular_identity(str(value)) |
| if identity["group_key"] != self.group_keys[index]: |
| raise ValueError("Reference SMILES and molecular identity disagree.") |
| if not np.isclose(np.float32(identity["exact_mass"]), self.masses[index], atol=5e-5, rtol=0): |
| raise ValueError("Reference exact masses differ from the frozen structure calculation.") |
| fingerprint = self.generator.GetFingerprint(identity["molecule"]) |
| features = np.zeros(2048, dtype=np.float32) |
| DataStructs.ConvertToNumpyArray(fingerprint, features) |
| if not np.array_equal(features, self.features[index]): |
| raise ValueError("Reference Morgan fingerprints differ from the current RDKit calculation.") |
| self.fingerprints.append(fingerprint) |
| self.connectivity.append(identity["smiles"]) |
| self.connectivity = np.asarray(self.connectivity, dtype=str) |
| self.training_keys = set(self.group_keys.tolist()) |
|
|
| def to(self, device): |
| requested = torch.device(device) |
| if requested.type not in {"cpu", "cuda"}: |
| raise ValueError("Only CPU and CUDA inference are supported.") |
| if requested.type == "cuda" and not torch.cuda.is_available(): |
| raise RuntimeError("CUDA was requested but is unavailable in the current worker.") |
| with self._lock: |
| if requested != self.device: |
| self.base.model.to(requested).eval() |
| self.base.device = requested |
| self.router.to(requested).eval() |
| self.device = requested |
| return self |
|
|
| def _neighbors(self, identity): |
| fingerprint = self.generator.GetFingerprint(identity["molecule"]) |
| features = np.zeros(2048, dtype=np.float32) |
| DataStructs.ConvertToNumpyArray(fingerprint, features) |
| similarities = np.asarray(DataStructs.BulkTanimotoSimilarity(fingerprint, self.fingerprints), dtype=np.float64) |
| eligible = np.flatnonzero(similarities >= PROTOCOL["minimum_similarity"]) |
| keep = [int(index) for index in eligible if self.group_keys[index] != identity["group_key"] |
| and self.connectivity[index] != identity["smiles"]] |
| keep.sort(key=lambda index: (-float(similarities[index]), str(self.group_keys[index]))) |
| indices = np.asarray(keep[:3], dtype=np.int64) |
| return indices, similarities[indices].astype(np.float32), features |
|
|
| def _transfer(self, identity, indices, similarities, query_features): |
| count = len(indices) |
| query = count |
| neighbors = np.full((count + 1, 3), -1, dtype=np.int64) |
| scores = np.zeros((count + 1, 3), dtype=np.float32) |
| neighbors[query, :count] = np.arange(count) |
| scores[query, :count] = similarities |
| data = { |
| "features": np.vstack((self.features[indices], query_features[None, :])), |
| "masses": np.concatenate((self.masses[indices], np.asarray([identity["exact_mass"]], dtype=np.float32))), |
| "spectra": np.vstack((self.spectra[indices], np.zeros((1, 751), dtype=np.float32))), |
| "neighbors": neighbors, |
| "similarities": scores, |
| } |
| transferred, covered = predict_router(self.router, data, np.asarray([query], dtype=np.int64), |
| device=str(self.device), mode="transport", batch_size=64, |
| min_similarity=0.35) |
| mz = np.arange(751) |
| query_mass = np.rint(np.float32(identity["exact_mass"])) |
| upper = min(int(query_mass) + 3, 750) |
| usable = [] |
| for index, similarity in zip(indices, similarities): |
| delta = int(query_mass) - int(np.rint(self.masses[index])) |
| shifted = mz + delta |
| positive = self.spectra[index] > 0 |
| copy_valid = (mz >= 1) & (mz <= upper) |
| shift_valid = (mz >= 1) & (shifted >= 1) & (shifted <= upper) & (shifted < 751) |
| usable.append(bool(similarity >= np.float32(0.35) and np.any(positive & (copy_valid | shift_valid)))) |
| return transferred, bool(covered[0]), usable |
|
|
| def predict(self, smiles, device="cpu"): |
| identity = molecular_identity(smiles) |
| with self._lock: |
| self.to(device) |
| indices, similarities, query_features = self._neighbors(identity) |
| masses = np.asarray([identity["exact_mass"]], dtype=np.float32) |
| base = np.asarray(self.base([identity["smiles"]], exact_mass=masses, batch_size=1), dtype=np.float32) |
| if base.shape != (1, 751) or not np.isfinite(base).all() or (base < 0).any() or base.sum() <= 0: |
| raise RuntimeError("The frozen base model returned an invalid or empty spectrum.") |
| base = l1_normalize(base) |
| extra, covered, usable = self._transfer(identity, indices, similarities, query_features) |
| extra = l1_normalize(extra) |
| spectrum = base.copy() |
| if covered: |
| spectrum = (1 - PROTOCOL["alpha"]) * base + PROTOCOL["alpha"] * extra |
| if not np.isfinite(spectrum).all() or (spectrum < 0).any() or spectrum.sum() <= 0: |
| raise RuntimeError("SMART returned an invalid or empty spectrum.") |
| warning_codes = [] |
| if identity["stereochemistry_ignored"]: |
| warning_codes.append("stereochemistry_ignored") |
| if not len(indices): |
| warning_codes.append("no_eligible_reference") |
| elif not covered: |
| warning_codes.append("reference_transfer_unavailable") |
| seen = identity["group_key"] in self.training_keys |
| if seen: |
| warning_codes.append("seen_in_training") |
| if identity["exact_mass"] > 750: |
| warning_codes.append("limited_mass_range") |
| relative = spectrum[0] / spectrum[0].max() * np.float32(100) |
| references = [{"group_key": str(self.group_keys[index]), "smiles": str(self.smiles[index]), |
| "similarity": float(similarity), "usable": valid} |
| for index, similarity, valid in zip(indices, similarities, usable)] |
| return { |
| "smiles": identity["smiles"], |
| "canonical_smiles": identity["smiles"], |
| "group_key": identity["group_key"], |
| "formula": identity["formula"], |
| "exact_mass": identity["exact_mass"], |
| "mz": np.arange(751, dtype=np.int64), |
| "spectrum": relative.astype(np.float32), |
| "intensity": relative.astype(np.float32), |
| "intensity_l1": spectrum[0].copy(), |
| "base_spectrum_l1": base[0].copy(), |
| "reference_count": int(sum(usable)), |
| "selected_reference_count": len(indices), |
| "reference_covered": covered, |
| "references": references, |
| "training_reference_match": seen, |
| "warnings": [WARNING_MESSAGES[code] for code in warning_codes], |
| "warning_codes": warning_codes, |
| "model_version": self.model_version, |
| "seed": 42, |
| "device": str(self.device), |
| "spectrum_type": "Predicted EI mass spectrum for GC–MS screening", |
| "intensity_scale": "Base peak = 100; m/z 0–750; nominal mass bins", |
| } |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--assets", type=Path, default=Path("assets")) |
| parser.add_argument("--smiles", required=True) |
| parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu") |
| parser.add_argument("--output", type=Path) |
| args = parser.parse_args() |
| result = SMARTPredictor(args.assets).predict(args.smiles, device=args.device) |
| serializable = {key: value.tolist() if isinstance(value, np.ndarray) else value for key, value in result.items()} |
| content = json.dumps(serializable, ensure_ascii=False, allow_nan=False, indent=2) |
| if args.output: |
| if args.output.exists(): |
| raise FileExistsError("Choose a new output path: " + str(args.output)) |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| args.output.write_text(content + "\n", encoding="utf-8") |
| else: |
| print(content) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
