Blame for inference.py

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()