bioassayalign_compatibility.py

from __future__ import annotations

import hashlib
import json
import os
import re
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path
from typing import Any

import numpy as np
import torch
import torch.nn.functional as F
from huggingface_hub import snapshot_download
from huggingface_hub.utils import disable_progress_bars
from rdkit import Chem, DataStructs, RDLogger
from rdkit.Chem import AllChem, Crippen, Descriptors, Lipinski, MACCSkeys, rdMolDescriptors
from rdkit.Chem.MolStandardize import rdMolStandardize
from sentence_transformers import SentenceTransformer
from torch import nn
from transformers.utils import logging as transformers_logging

os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
disable_progress_bars()
transformers_logging.set_verbosity_error()
RDLogger.DisableLog("rdApp.*")

DEFAULT_ASSAY_TASK = (
    "Given a bioassay description and metadata, represent the assay for ranking compatible small molecules."
)
SECTION_ORDER = [
    "ASSAY_TITLE",
    "DESCRIPTION",
    "ORGANISM",
    "READOUT",
    "ASSAY_FORMAT",
    "ASSAY_TYPE",
    "TARGET_UNIPROT",
]
ASSAY_SECTION_RE = re.compile(
    r"\[(ASSAY_TITLE|DESCRIPTION|ORGANISM|READOUT|ASSAY_FORMAT|ASSAY_TYPE|TARGET_UNIPROT)\]\n"
)
ORGANISM_ALIASES = {
    "9606": "homo_sapiens",
    "10090": "mus_musculus",
    "10116": "rattus_norvegicus",
    "4932": "saccharomyces_cerevisiae",
}
DEFAULT_DESCRIPTOR_NAMES = (
    "mol_wt",
    "logp",
    "tpsa",
    "heavy_atoms",
    "hbd",
    "hba",
    "rot_bonds",
    "ring_count",
    "aromatic_rings",
    "aliphatic_rings",
    "saturated_rings",
    "fraction_csp3",
    "heteroatoms",
    "amide_bonds",
    "fragments",
    "formal_charge",
    "max_atomic_num",
    "metal_atom_count",
    "halogen_count",
    "nitrogen_count",
    "oxygen_count",
    "sulfur_count",
    "phosphorus_count",
    "fluorine_count",
    "chlorine_count",
    "bromine_count",
    "iodine_count",
    "aromatic_atom_count",
    "spiro_atoms",
    "bridgehead_atoms",
)
ORGANIC_LIKE_ATOMIC_NUMBERS = {1, 5, 6, 7, 8, 9, 14, 15, 16, 17, 35, 53}


@dataclass
class AssayQuery:
    title: str = ""
    description: str = ""
    organism: str = ""
    readout: str = ""
    assay_format: str = ""
    assay_type: str = ""
    target_uniprot: list[str] | None = None


def serialize_assay_query(query: AssayQuery) -> str:
    targets = ", ".join(query.target_uniprot or [])
    values = {
        "ASSAY_TITLE": query.title.strip(),
        "DESCRIPTION": query.description.strip(),
        "ORGANISM": query.organism.strip(),
        "READOUT": query.readout.strip(),
        "ASSAY_FORMAT": query.assay_format.strip(),
        "ASSAY_TYPE": query.assay_type.strip(),
        "TARGET_UNIPROT": targets.strip(),
    }
    return "\n\n".join(f"[{key}]\n{values[key]}" for key in SECTION_ORDER)


def _format_assay_query(assay_text: str, task_description: str) -> str:
    return f"Instruct: {task_description.strip()}\nQuery: {assay_text.strip()}"


def _parse_assay_sections(assay_text: str) -> dict[str, str]:
    sections = {key: "" for key in SECTION_ORDER}
    parts = ASSAY_SECTION_RE.split(assay_text)
    for idx in range(1, len(parts), 2):
        key = parts[idx]
        value = parts[idx + 1] if idx + 1 < len(parts) else ""
        if key in sections:
            sections[key] = value.strip()
    return sections


def _normalize_metadata_token(value: str) -> str:
    return re.sub(r"[^a-z0-9]+", "_", value.lower()).strip("_")


def _normalize_organism_token(value: str) -> str:
    raw = value.strip()
    if not raw:
        return ""
    aliased = ORGANISM_ALIASES.get(raw, raw)
    return _normalize_metadata_token(aliased)


def _hash_bucket(value: str, dim: int) -> int:
    return abs(hash(value)) % max(dim, 1)


def _assay_metadata_vector(assay_text: str, *, dim: int) -> np.ndarray:
    if dim <= 0:
        return np.zeros((0,), dtype=np.float32)
    sections = _parse_assay_sections(assay_text)
    tokens: list[str] = []
    organism = _normalize_organism_token(sections.get("ORGANISM", ""))
    if organism:
        tokens.append(f"organism:{organism}")
    for key in ("READOUT", "ASSAY_FORMAT", "ASSAY_TYPE"):
        value = _normalize_metadata_token(sections.get(key, ""))
        if value:
            tokens.append(f"{key.lower()}:{value}")
    for target in sections.get("TARGET_UNIPROT", "").split(","):
        token = target.strip().upper()
        if token:
            tokens.append(f"target:{token}")
    vec = np.zeros((dim,), dtype=np.float32)
    for token in tokens:
        vec[_hash_bucket(token, dim)] += 1.0
    norm = float(np.linalg.norm(vec))
    if norm > 0:
        vec /= norm
    return vec


@lru_cache(maxsize=1_000_000)
def _standardize_smiles_v2_cached(smiles: str) -> str | None:
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return None
    try:
        mol = rdMolStandardize.Cleanup(mol)
        mol = rdMolStandardize.FragmentParent(mol)
        mol = rdMolStandardize.Uncharger().uncharge(mol)
        mol = rdMolStandardize.TautomerEnumerator().Canonicalize(mol)
        Chem.SanitizeMol(mol)
    except Exception:
        return None
    if mol.GetNumHeavyAtoms() < 2:
        return None
    standardized = Chem.MolToSmiles(mol, canonical=True, isomericSmiles=True)
    if not standardized or "." in standardized:
        return None
    return standardized


def standardize_smiles_v2(smiles: str | None) -> str | None:
    if not smiles:
        return None
    token = smiles.strip()
    if not token:
        return None
    return _standardize_smiles_v2_cached(token)


def smiles_sha256(smiles: str) -> str:
    return hashlib.sha256(smiles.encode("utf-8")).hexdigest()


def _count_atomic_nums(mol) -> dict[int, int]:
    counts: dict[int, int] = {}
    for atom in mol.GetAtoms():
        atomic_num = int(atom.GetAtomicNum())
        counts[atomic_num] = counts.get(atomic_num, 0) + 1
    return counts


def _morgan_bits_from_mol(mol, *, radius: int, n_bits: int, use_chirality: bool) -> np.ndarray:
    fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius, nBits=n_bits, useChirality=use_chirality)
    arr = np.zeros((n_bits,), dtype=np.uint8)
    DataStructs.ConvertToNumpyArray(fp, arr)
    return arr


def _maccs_bits_from_mol(mol) -> np.ndarray:
    fp = MACCSkeys.GenMACCSKeys(mol)
    arr = np.zeros((fp.GetNumBits(),), dtype=np.uint8)
    DataStructs.ConvertToNumpyArray(fp, arr)
    return arr


def _molecule_descriptor_vector(mol, *, names: tuple[str, ...] = DEFAULT_DESCRIPTOR_NAMES) -> np.ndarray:
    counts = _count_atomic_nums(mol)
    fragments = Chem.GetMolFrags(mol)
    formal_charge = sum(int(atom.GetFormalCharge()) for atom in mol.GetAtoms())
    max_atomic_num = max(counts) if counts else 0
    metal_atom_count = sum(
        count for atomic_num, count in counts.items() if atomic_num not in ORGANIC_LIKE_ATOMIC_NUMBERS
    )
    halogen_count = sum(counts.get(item, 0) for item in (9, 17, 35, 53))
    aromatic_atom_count = sum(1 for atom in mol.GetAtoms() if atom.GetIsAromatic())
    values = {
        "mol_wt": float(Descriptors.MolWt(mol)),
        "logp": float(Crippen.MolLogP(mol)),
        "tpsa": float(rdMolDescriptors.CalcTPSA(mol)),
        "heavy_atoms": float(mol.GetNumHeavyAtoms()),
        "hbd": float(Lipinski.NumHDonors(mol)),
        "hba": float(Lipinski.NumHAcceptors(mol)),
        "rot_bonds": float(Lipinski.NumRotatableBonds(mol)),
        "ring_count": float(rdMolDescriptors.CalcNumRings(mol)),
        "aromatic_rings": float(rdMolDescriptors.CalcNumAromaticRings(mol)),
        "aliphatic_rings": float(rdMolDescriptors.CalcNumAliphaticRings(mol)),
        "saturated_rings": float(rdMolDescriptors.CalcNumSaturatedRings(mol)),
        "fraction_csp3": float(rdMolDescriptors.CalcFractionCSP3(mol)),
        "heteroatoms": float(rdMolDescriptors.CalcNumHeteroatoms(mol)),
        "amide_bonds": float(rdMolDescriptors.CalcNumAmideBonds(mol)),
        "fragments": float(len(fragments)),
        "formal_charge": float(formal_charge),
        "max_atomic_num": float(max_atomic_num),
        "metal_atom_count": float(metal_atom_count),
        "halogen_count": float(halogen_count),
        "nitrogen_count": float(counts.get(7, 0)),
        "oxygen_count": float(counts.get(8, 0)),
        "sulfur_count": float(counts.get(16, 0)),
        "phosphorus_count": float(counts.get(15, 0)),
        "fluorine_count": float(counts.get(9, 0)),
        "chlorine_count": float(counts.get(17, 0)),
        "bromine_count": float(counts.get(35, 0)),
        "iodine_count": float(counts.get(53, 0)),
        "aromatic_atom_count": float(aromatic_atom_count),
        "spiro_atoms": float(rdMolDescriptors.CalcNumSpiroAtoms(mol)),
        "bridgehead_atoms": float(rdMolDescriptors.CalcNumBridgeheadAtoms(mol)),
    }
    return np.asarray([values[name] for name in names], dtype=np.float32)


class CompatibilityHead(nn.Module):
    def __init__(self, *, assay_dim: int, molecule_dim: int, projection_dim: int, hidden_dim: int, dropout: float) -> None:
        super().__init__()
        self.assay_norm = nn.LayerNorm(assay_dim)
        self.assay_proj = nn.Linear(assay_dim, projection_dim)
        self.mol_norm = nn.LayerNorm(molecule_dim)
        self.mol_proj = nn.Linear(molecule_dim, projection_dim, bias=False)
        self.score_mlp = nn.Sequential(
            nn.Linear(projection_dim * 4, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, 1),
        )
        self.dot_scale = nn.Parameter(torch.tensor(1.0, dtype=torch.float32))

    def encode_assay(self, assay_features: torch.Tensor) -> torch.Tensor:
        vec = self.assay_proj(self.assay_norm(assay_features))
        return F.normalize(vec, p=2, dim=-1)

    def encode_molecule(self, molecule_features: torch.Tensor) -> torch.Tensor:
        vec = self.mol_proj(self.mol_norm(molecule_features))
        return F.normalize(vec, p=2, dim=-1)

    def score_candidates(self, assay_features: torch.Tensor, candidate_features: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        assay_vec = self.encode_assay(assay_features)
        mol_vec = self.encode_molecule(candidate_features)
        assay_expand = assay_vec.unsqueeze(1).expand(-1, mol_vec.shape[1], -1)
        dot_scores = (assay_expand * mol_vec).sum(dim=-1)
        mlp_input = torch.cat(
            [
                assay_expand,
                mol_vec,
                assay_expand * mol_vec,
                torch.abs(assay_expand - mol_vec),
            ],
            dim=-1,
        )
        mlp_scores = self.score_mlp(mlp_input).squeeze(-1)
        logits = dot_scores * self.dot_scale + mlp_scores
        return logits, assay_vec, mol_vec

    def score_pairs(self, assay_features: torch.Tensor, molecule_features: torch.Tensor) -> torch.Tensor:
        assay_vec = self.encode_assay(assay_features)
        mol_vec = self.encode_molecule(molecule_features)
        dot_scores = (assay_vec * mol_vec).sum(dim=-1)
        mlp_input = torch.cat(
            [assay_vec, mol_vec, assay_vec * mol_vec, torch.abs(assay_vec - mol_vec)],
            dim=-1,
        )
        mlp_scores = self.score_mlp(mlp_input).squeeze(-1)
        return dot_scores * self.dot_scale + mlp_scores


class BioAssayAlignCompatibilityModel:
    def __init__(
        self,
        assay_encoder: SentenceTransformer,
        compatibility_head: CompatibilityHead,
        *,
        assay_task_description: str,
        fingerprint_radii: tuple[int, ...],
        fingerprint_bits: int,
        use_chirality: bool,
        use_maccs: bool,
        use_rdkit_descriptors: bool,
        descriptor_names: tuple[str, ...],
        descriptor_mean: np.ndarray | None,
        descriptor_std: np.ndarray | None,
        use_assay_metadata_features: bool,
        assay_metadata_dim: int,
    ) -> None:
        self.assay_encoder = assay_encoder
        self.compatibility_head = compatibility_head.eval()
        self.assay_task_description = assay_task_description
        self.fingerprint_radii = fingerprint_radii
        self.fingerprint_bits = fingerprint_bits
        self.use_chirality = use_chirality
        self.use_maccs = use_maccs
        self.use_rdkit_descriptors = use_rdkit_descriptors
        self.descriptor_names = descriptor_names
        self.descriptor_mean = descriptor_mean
        self.descriptor_std = descriptor_std
        self.use_assay_metadata_features = use_assay_metadata_features
        self.assay_metadata_dim = assay_metadata_dim

    def _build_assay_feature_array(self, assay_text: str) -> np.ndarray:
        query = _format_assay_query(assay_text, self.assay_task_description)
        assay_features = self.assay_encoder.encode(
            [query],
            batch_size=1,
            normalize_embeddings=True,
            show_progress_bar=False,
            convert_to_numpy=True,
        )[0].astype(np.float32)
        if self.use_assay_metadata_features and self.assay_metadata_dim > 0:
            metadata_vec = _assay_metadata_vector(assay_text, dim=self.assay_metadata_dim)
            assay_features = np.concatenate([assay_features, metadata_vec.astype(np.float32)], axis=0)
        return assay_features

    def build_molecule_feature_matrix(self, smiles_values: list[str]) -> np.ndarray:
        rows: list[np.ndarray] = []
        for smiles in smiles_values:
            rows.append(
                _smiles_to_molecule_features(
                    smiles,
                    radii=self.fingerprint_radii,
                    n_bits=self.fingerprint_bits,
                    use_chirality=self.use_chirality,
                    use_maccs=self.use_maccs,
                    use_rdkit_descriptors=self.use_rdkit_descriptors,
                    descriptor_names=self.descriptor_names,
                    descriptor_mean=self.descriptor_mean,
                    descriptor_std=self.descriptor_std,
                )
            )
        return np.stack(rows, axis=0).astype(np.float32)

    def score(self, assay_text: str, smiles: str) -> float:
        assay_features = self._build_assay_feature_array(assay_text)
        molecule_features = self.build_molecule_feature_matrix([smiles])[0]
        assay_tensor = torch.from_numpy(assay_features).unsqueeze(0)
        molecule_tensor = torch.from_numpy(molecule_features).unsqueeze(0)
        with torch.no_grad():
            score = self.compatibility_head.score_pairs(assay_tensor, molecule_tensor)
        return float(score.item())


def _load_sentence_transformer(model_name: str):
    dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
    encoder = SentenceTransformer(
        model_name,
        trust_remote_code=True,
        model_kwargs={"torch_dtype": dtype},
    )
    if getattr(encoder, "tokenizer", None) is not None:
        encoder.tokenizer.padding_side = "left"
    return encoder


def _load_feature_spec(cfg: dict[str, Any], metadata: dict[str, Any], checkpoint: dict[str, Any]) -> dict[str, Any]:
    spec = checkpoint.get("molecule_feature_spec") or metadata.get("molecule_feature_spec")
    if spec:
        return spec
    radii = tuple(int(item) for item in (cfg.get("fingerprint_radii") or [cfg.get("fingerprint_radius", 2)]))
    return {
        "fingerprint_radii": list(radii),
        "fingerprint_bits": int(cfg["fingerprint_bits"]),
        "use_chirality": bool(cfg.get("use_chirality", False)),
        "use_maccs": bool(cfg.get("use_maccs", False)),
        "use_rdkit_descriptors": bool(cfg.get("use_rdkit_descriptors", False)),
        "descriptor_names": [],
        "descriptor_mean": None,
        "descriptor_std": None,
    }


def _smiles_to_molecule_features(
    smiles: str,
    *,
    radii: tuple[int, ...],
    n_bits: int,
    use_chirality: bool,
    use_maccs: bool,
    use_rdkit_descriptors: bool,
    descriptor_names: tuple[str, ...],
    descriptor_mean: np.ndarray | None,
    descriptor_std: np.ndarray | None,
) -> np.ndarray:
    normalized = standardize_smiles_v2(smiles) or smiles
    mol = Chem.MolFromSmiles(normalized)
    if mol is None:
        raise ValueError(f"Could not parse SMILES: {normalized}")
    bit_blocks: list[np.ndarray] = [
        _morgan_bits_from_mol(mol, radius=int(radius), n_bits=n_bits, use_chirality=use_chirality)
        for radius in radii
    ]
    if use_maccs:
        bit_blocks.append(_maccs_bits_from_mol(mol))
    output_blocks: list[np.ndarray] = [np.concatenate(bit_blocks, axis=0).astype(np.float32)]
    if use_rdkit_descriptors and descriptor_names:
        dense = _molecule_descriptor_vector(mol, names=descriptor_names)
        if descriptor_mean is not None and descriptor_std is not None:
            dense = (dense - descriptor_mean) / descriptor_std
        output_blocks.append(dense.astype(np.float32))
    return np.concatenate(output_blocks, axis=0).astype(np.float32)


def load_compatibility_model(model_dir: str | Path) -> BioAssayAlignCompatibilityModel:
    model_path = Path(model_dir)
    checkpoint = torch.load(model_path / "best_model.pt", map_location="cpu", weights_only=False)
    metadata = json.loads((model_path / "training_metadata.json").read_text())
    cfg = metadata["config"]
    feature_spec = _load_feature_spec(cfg, metadata, checkpoint)

    encoder = _load_sentence_transformer(checkpoint.get("assay_model_name") or cfg["assay_model_name"])
    assay_dim = int(checkpoint["model_state_dict"]["assay_proj.weight"].shape[1])
    molecule_dim = int(checkpoint["model_state_dict"]["mol_proj.weight"].shape[1])
    head = CompatibilityHead(
        assay_dim=assay_dim,
        molecule_dim=molecule_dim,
        projection_dim=int(cfg["projection_dim"]),
        hidden_dim=int(cfg["hidden_dim"]),
        dropout=float(cfg["dropout"]),
    )
    load_result = head.load_state_dict(checkpoint["model_state_dict"], strict=False)
    allowed_missing = {"mol_norm.weight", "mol_norm.bias"}
    unexpected = set(load_result.unexpected_keys)
    missing = set(load_result.missing_keys)
    if unexpected or (missing - allowed_missing):
        raise RuntimeError(
            "Compatibility checkpoint load mismatch: "
            f"unexpected={sorted(unexpected)} missing={sorted(missing)}"
        )
    return BioAssayAlignCompatibilityModel(
        assay_encoder=encoder,
        compatibility_head=head,
        assay_task_description=checkpoint.get("assay_task_description") or cfg["assay_task_description"],
        fingerprint_radii=tuple(int(item) for item in feature_spec.get("fingerprint_radii") or [2]),
        fingerprint_bits=int(feature_spec.get("fingerprint_bits", cfg.get("fingerprint_bits", 2048))),
        use_chirality=bool(feature_spec.get("use_chirality", cfg.get("use_chirality", False))),
        use_maccs=bool(feature_spec.get("use_maccs", False)),
        use_rdkit_descriptors=bool(feature_spec.get("use_rdkit_descriptors", False)),
        descriptor_names=tuple(feature_spec.get("descriptor_names") or ()),
        descriptor_mean=np.array(feature_spec["descriptor_mean"], dtype=np.float32)
        if feature_spec.get("descriptor_mean") is not None
        else None,
        descriptor_std=np.array(feature_spec["descriptor_std"], dtype=np.float32)
        if feature_spec.get("descriptor_std") is not None
        else None,
        use_assay_metadata_features=bool(cfg.get("use_assay_metadata_features", False)),
        assay_metadata_dim=int(cfg.get("assay_metadata_dim", 0) or 0),
    )


def load_compatibility_model_from_hub(repo_id: str) -> BioAssayAlignCompatibilityModel:
    snapshot_path = snapshot_download(
        repo_id=repo_id,
        repo_type="model",
        allow_patterns=["best_model.pt", "training_metadata.json"],
    )
    return load_compatibility_model(snapshot_path)


def rank_compounds(
    model: BioAssayAlignCompatibilityModel,
    *,
    assay_text: str,
    smiles_list: list[str],
    top_k: int | None = None,
) -> list[dict[str, Any]]:
    if not smiles_list:
        return []

    assay_features = model._build_assay_feature_array(assay_text)
    assay_tensor = torch.from_numpy(assay_features.astype(np.float32)).unsqueeze(0)

    valid_items: list[tuple[str, str]] = []
    invalid_items: list[dict[str, Any]] = []
    for raw_smiles in smiles_list:
        standardized = standardize_smiles_v2(raw_smiles)
        if standardized is None:
            invalid_items.append(
                {
                    "input_smiles": raw_smiles,
                    "canonical_smiles": None,
                    "smiles_hash": None,
                    "score": None,
                    "valid": False,
                    "error": "invalid_smiles",
                }
            )
            continue
        valid_items.append((raw_smiles, standardized))

    ranked_items: list[dict[str, Any]] = []
    if valid_items:
        feature_matrix = model.build_molecule_feature_matrix([item[1] for item in valid_items])
        candidate_tensor = torch.from_numpy(feature_matrix).unsqueeze(0)
        with torch.no_grad():
            logits, _, _ = model.compatibility_head.score_candidates(
                assay_tensor.to(dtype=torch.float32),
                candidate_tensor.to(dtype=torch.float32),
            )
        scores = logits.squeeze(0).cpu().numpy().tolist()
        for (raw_smiles, canonical), score in zip(valid_items, scores, strict=True):
            ranked_items.append(
                {
                    "input_smiles": raw_smiles,
                    "canonical_smiles": canonical,
                    "smiles_hash": smiles_sha256(canonical),
                    "score": float(score),
                    "valid": True,
                }
            )
        ranked_items.sort(key=lambda item: item["score"], reverse=True)
        if top_k is not None and top_k > 0:
            ranked_items = ranked_items[:top_k]

    return ranked_items + invalid_items


def list_softmax_scores(scores: list[float], temperature: float = 1.0) -> list[float]:
    values = np.asarray(scores, dtype=np.float32) / max(float(temperature), 1e-6)
    values = values - values.max()
    probs = np.exp(values)
    probs = probs / probs.sum()
    return probs.tolist()