Add self-contained inference helper and clarify score semantics

README.md

@@ -57,6 +57,12 @@ This artifact uses:
 
 The final score comes from the learned compatibility head. It is not just a raw embedding dot product.
 
+This repo also includes a self-contained inference helper:
+- `bioassayalign_compatibility.py`
+- `requirements.txt`
+
+You do not need the project GitHub repo to run the published model.
+
 ## Training Data
 
 Training uses a frozen public assay-compound corpus derived from:
@@ -137,39 +143,97 @@ Candidate list ranked by the model:
 
 The raw values above are model scores. In practice, read them as list-relative ranking values, not calibrated probabilities.
 
+## What The Raw Score Means
+
+The raw output is a learned ranking score from the compatibility head.
+
+You can think of it as a **logit-like utility value**:
+- higher is better
+- differences inside the same submitted list matter
+- absolute values across unrelated lists are not directly comparable
+
+Example:
+- a top candidate with score `-4.7`
+- another candidate with score `-20.0`
+
+does **not** mean the first compound has negative biological value. It only means the first item scored much better than the second one for that submitted assay-and-list context.
+
+If you want a normalized shortlist view for one submitted list, you can convert the raw scores with:
+- a min-max `0–100` relative ranking scale, or
+- a softmax over the submitted list
+
+Softmax example for one list:
+
+```python
+from bioassayalign_compatibility import list_softmax_scores
+
+scores = [-4.6947, -15.0503, -20.0474]
+relative_probs = list_softmax_scores(scores)
+print(relative_probs)
+```
+
+Important:
+- those softmax values are **not calibrated probabilities of success**
+- they are only a normalized way to compare the candidates you submitted together
+
 ## How To Run It Locally
 
 ### Minimal local check from this repo
 
-This downloads only the model weights and metadata, not the raw assay dataset.
+This downloads only the published model files from this repo, not the raw assay dataset.
 
 ```bash
-MODEL_REPO_ID='lighteternal/BioAssayAlign-Qwen3-Embedding-0.6B-Compatibility' \
-LOCAL_MODEL_DIR='data/hf_compat_model_check' \
-bash scripts/score_compatibility_from_hf.sh \
-  --assay-title 'JAK2 inhibition assay' \
-  --description 'Cell-based luminescence assay measuring JAK2 inhibition in HEK293 cells.' \
-  --organism 'Homo sapiens' \
-  --readout 'luminescence' \
-  --assay-format 'cell-based' \
-  --assay-type 'inhibition' \
-  --target-uniprot O60674 \
-  --smiles 'CC(=O)Nc1ncc(C#N)c(Nc2ccc(F)c(Cl)c2)n1' \
-  --smiles 'c1ccccc1' \
-  --smiles 'CCO'
+python -m pip install -r requirements.txt
+python - <<'PY'
+from bioassayalign_compatibility import (
+    AssayQuery,
+    load_compatibility_model_from_hub,
+    rank_compounds,
+    serialize_assay_query,
+)
+
+model = load_compatibility_model_from_hub(
+    "lighteternal/BioAssayAlign-Qwen3-Embedding-0.6B-Compatibility"
+)
+assay_text = serialize_assay_query(
+    AssayQuery(
+        title="JAK2 inhibition assay",
+        description="Cell-based luminescence assay measuring JAK2 inhibition in HEK293 cells.",
+        organism="Homo sapiens",
+        readout="luminescence",
+        assay_format="cell-based",
+        assay_type="inhibition",
+        target_uniprot=["O60674"],
+    )
+)
+
+results = rank_compounds(
+    model,
+    assay_text=assay_text,
+    smiles_list=[
+        "CC(=O)Nc1ncc(C#N)c(Nc2ccc(F)c(Cl)c2)n1",
+        "c1ccccc1",
+        "CCO",
+    ],
+)
+for row in results:
+    print(row)
+PY
 ```
 
 ### Python usage
 
 ```python
-from bioassayalign.compat_inference import (
+from bioassayalign_compatibility import (
     AssayQuery,
-    load_compatibility_model,
+    load_compatibility_model_from_hub,
     rank_compounds,
     serialize_assay_query,
 )
 
-model = load_compatibility_model("/path/to/model_dir")
+model = load_compatibility_model_from_hub(
+    "lighteternal/BioAssayAlign-Qwen3-Embedding-0.6B-Compatibility"
+)
 assay_text = serialize_assay_query(
     AssayQuery(
         title="JAK2 inhibition assay",

bioassayalign_compatibility.py

@@ -0,0 +1,439 @@
+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 _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,
+        metadata_dim: int = 0,
+    ) -> None:
+        super().__init__()
+        self.metadata_dim = metadata_dim
+        assay_input_dim = assay_dim + metadata_dim
+        self.assay_proj = nn.Sequential(
+            nn.Linear(assay_input_dim, projection_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+        )
+        self.molecule_proj = nn.Sequential(
+            nn.Linear(molecule_dim, projection_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+        )
+        self.scorer = nn.Sequential(
+            nn.Linear(projection_dim * 4, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, 1),
+        )
+
+    def forward(self, assay_vec: torch.Tensor, molecule_vec: torch.Tensor, assay_metadata: torch.Tensor | None = None) -> torch.Tensor:
+        if assay_metadata is not None and assay_metadata.numel():
+            assay_input = torch.cat([assay_vec, assay_metadata], dim=-1)
+        else:
+            assay_input = assay_vec
+        assay_hidden = self.assay_proj(assay_input)
+        molecule_hidden = self.molecule_proj(molecule_vec)
+        interaction = torch.cat(
+            [
+                assay_hidden,
+                molecule_hidden,
+                assay_hidden * molecule_hidden,
+                torch.abs(assay_hidden - molecule_hidden),
+            ],
+            dim=-1,
+        )
+        return self.scorer(interaction).squeeze(-1)
+
+
+class CompatibilityModel:
+    def __init__(self, assay_encoder: SentenceTransformer, metadata: dict[str, Any], head_state: dict[str, Any], *, device: str | None = None) -> None:
+        self.metadata = metadata
+        self.config = metadata["config"]
+        self.feature_spec = metadata["molecule_feature_spec"]
+        self.metadata_dim = int(self.config.get("assay_metadata_dim", 0))
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.assay_encoder = assay_encoder
+        self.assay_encoder.max_seq_length = 512
+        self.assay_dim = int(self.assay_encoder.get_sentence_embedding_dimension())
+        self.molecule_dim = int(metadata["feature_counts"]["molecule_dim"])
+        self.head = CompatibilityHead(
+            assay_dim=self.assay_dim,
+            molecule_dim=self.molecule_dim,
+            projection_dim=int(self.config["projection_dim"]),
+            hidden_dim=int(self.config["hidden_dim"]),
+            dropout=float(self.config["dropout"]),
+            metadata_dim=self.metadata_dim,
+        ).to(self.device)
+        self.head.load_state_dict(head_state)
+        self.head.eval()
+
+    def encode_assay(self, assay_text: str) -> tuple[torch.Tensor, torch.Tensor | None]:
+        embedding = self.assay_encoder.encode(
+            [assay_text],
+            convert_to_numpy=True,
+            show_progress_bar=False,
+            normalize_embeddings=True,
+            prompt_name="query",
+            prompt=DEFAULT_ASSAY_TASK,
+        )[0].astype(np.float32)
+        assay_vec = torch.from_numpy(embedding).unsqueeze(0).to(self.device)
+        metadata_vec = _assay_metadata_vector(assay_text, dim=self.metadata_dim)
+        metadata_tensor = None
+        if metadata_vec.size:
+            metadata_tensor = torch.from_numpy(metadata_vec).unsqueeze(0).to(self.device)
+        return assay_vec, metadata_tensor
+
+    def score_feature_matrix(self, assay_text: str, feature_matrix: np.ndarray) -> np.ndarray:
+        assay_vec, metadata_tensor = self.encode_assay(assay_text)
+        molecule_tensor = torch.from_numpy(feature_matrix).to(self.device)
+        with torch.inference_mode():
+            assay_repeat = assay_vec.repeat(molecule_tensor.size(0), 1)
+            metadata_repeat = metadata_tensor.repeat(molecule_tensor.size(0), 1) if metadata_tensor is not None else None
+            scores = self.head(assay_repeat, molecule_tensor, metadata_repeat)
+        return scores.detach().cpu().numpy()
+
+
+def build_molecule_feature_vector(smiles: str, feature_spec: dict[str, Any]) -> np.ndarray | None:
+    standardized = standardize_smiles_v2(smiles)
+    if standardized is None:
+        return None
+    mol = Chem.MolFromSmiles(standardized)
+    if mol is None:
+        return None
+    parts: list[np.ndarray] = []
+    for radius in feature_spec.get("fingerprint_radii", [2, 3]):
+        parts.append(
+            _morgan_bits_from_mol(
+                mol,
+                radius=int(radius),
+                n_bits=int(feature_spec.get("fingerprint_bits", 2048)),
+                use_chirality=bool(feature_spec.get("use_chirality", True)),
+            ).astype(np.float32)
+        )
+    if feature_spec.get("use_maccs", True):
+        parts.append(_maccs_bits_from_mol(mol).astype(np.float32))
+    if feature_spec.get("use_rdkit_descriptors", True):
+        descriptor_values = _molecule_descriptor_vector(
+            mol,
+            names=tuple(feature_spec.get("descriptor_names", DEFAULT_DESCRIPTOR_NAMES)),
+        )
+        descriptor_mean = np.asarray(feature_spec["descriptor_mean"], dtype=np.float32)
+        descriptor_std = np.asarray(feature_spec["descriptor_std"], dtype=np.float32)
+        parts.append(((descriptor_values - descriptor_mean) / (descriptor_std + 1e-6)).astype(np.float32))
+    if not parts:
+        return None
+    return np.concatenate(parts, axis=0).astype(np.float32)
+
+
+def load_compatibility_model(model_dir: str | Path, *, device: str | None = None) -> CompatibilityModel:
+    model_path = Path(model_dir)
+    training_metadata = json.loads((model_path / "training_metadata.json").read_text())
+    checkpoint = torch.load(model_path / "best_model.pt", map_location="cpu")
+    assay_model_name = training_metadata["config"]["assay_model_name"]
+    assay_encoder = SentenceTransformer(assay_model_name, device=device or ("cuda" if torch.cuda.is_available() else "cpu"))
+    return CompatibilityModel(assay_encoder, training_metadata, checkpoint["head_state"], device=device)
+
+
+def load_compatibility_model_from_hub(repo_id: str, *, device: str | None = None) -> CompatibilityModel:
+    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, device=device)
+
+
+def rank_compounds(model: CompatibilityModel, assay_text: str, smiles_list: list[str], *, top_k: int | None = None) -> list[dict[str, Any]]:
+    valid_inputs: list[tuple[str, str, np.ndarray]] = []
+    invalid_rows: list[dict[str, Any]] = []
+    for item in smiles_list:
+        feature_vec = build_molecule_feature_vector(item, model.feature_spec)
+        standardized = standardize_smiles_v2(item)
+        if feature_vec is None or standardized is None:
+            invalid_rows.append({"input_smiles": item, "valid": False, "error": "invalid_smiles"})
+            continue
+        valid_inputs.append((item, standardized, feature_vec))
+    valid_rows: list[dict[str, Any]] = []
+    if valid_inputs:
+        feature_matrix = np.stack([entry[2] for entry in valid_inputs], axis=0).astype(np.float32)
+        scores = model.score_feature_matrix(assay_text, feature_matrix)
+        for (input_smiles, standardized, _), score in zip(valid_inputs, scores):
+            valid_rows.append(
+                {
+                    "input_smiles": input_smiles,
+                    "canonical_smiles": standardized,
+                    "smiles_hash": smiles_sha256(standardized),
+                    "score": float(score),
+                    "valid": True,
+                }
+            )
+        valid_rows.sort(key=lambda item: item["score"], reverse=True)
+        if top_k:
+            valid_rows = valid_rows[:top_k]
+    return valid_rows + invalid_rows
+
+
+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()

requirements.txt

@@ -0,0 +1,6 @@
+torch>=2.2
+transformers>=4.50
+sentence-transformers>=3.0
+huggingface_hub>=0.30
+numpy>=1.26
+rdkit>=2023.9