Upload BioAssayAlign compatibility Space bundle

README.md

@@ -1,12 +1,77 @@
 ---
 title: BioAssayAlign Compatibility Explorer
-emoji: 🦀
-colorFrom: gray
-colorTo: yellow
+emoji: 🧪
+colorFrom: blue
+colorTo: gray
 sdk: gradio
-sdk_version: 6.9.0
 app_file: app.py
 pinned: false
+license: mit
+short_description: Rank candidate molecules for a bioassay.
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# BioAssayAlign Compatibility Explorer
+
+This Space is a scientist-facing demo for **assay-conditioned compound ranking**.
+
+You provide:
+- a bioassay description and optional metadata
+- a list of candidate SMILES
+
+The model returns:
+- a ranked list of molecules
+- a compatibility score for each one
+- explicit flags for invalid SMILES
+
+## What It Is
+
+This is not a chatbot and it is not a potency predictor.
+
+It is a **ranking model** trained on a frozen public bioassay dataset built from PubChem BioAssay and ChEMBL. It is designed to answer:
+
+> “Given this assay, which molecules should I look at first?”
+
+## What The Score Means
+
+- Higher score = the model believes the molecule is more compatible with the assay than lower-ranked candidates in the same list.
+- The score is **not** a probability.
+- The score is best used for **ranking**, not absolute decision thresholds.
+
+## Recommended Input Style
+
+The model is most reliable when assay information is provided as structured fields:
+- title
+- description
+- organism
+- readout
+- assay format
+- assay type
+- target UniProt IDs
+
+You can paste SMILES directly or upload a CSV with a `smiles` or `canonical_smiles` column.
+
+## Good Uses
+
+- ranking a screening shortlist for a new assay concept
+- triaging compounds before a more expensive downstream model or wet-lab step
+- testing how sensitive rankings are to assay wording and metadata
+
+## Limits
+
+- This is a public-data model, not a medicinal chemistry oracle.
+- It does not predict IC50 directly.
+- It is strongest as a **relative ranking tool** over a candidate list you already care about.
+
+## Runtime Notes
+
+- The first request can be slower because the Space has to load the model.
+- Large candidate lists increase runtime. For interactive use, start with a few hundred molecules.
+
+## Model
+
+The Space reads the model repo from the `MODEL_REPO_ID` environment variable.
+
+Default:
+- `lighteternal/BioAssayAlign-Qwen3-Embedding-0.6B-Compatibility`
+
+If the champion changes later, the Space can point to a new model repo without changing the UI.

app.py

@@ -0,0 +1,407 @@
+from __future__ import annotations
+
+import csv
+import os
+import tempfile
+from pathlib import Path
+from typing import Any
+
+import gradio as gr
+import pandas as pd
+
+from space_runtime import AssayQuery, load_compatibility_model_from_hub, rank_compounds, serialize_assay_query
+
+MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", "lighteternal/BioAssayAlign-Qwen3-Embedding-0.6B-Compatibility")
+MAX_INPUT_SMILES = int(os.getenv("MAX_INPUT_SMILES", "3000"))
+DEFAULT_TOP_K = int(os.getenv("DEFAULT_TOP_K", "50"))
+
+CSS = """
+@import url('https://fonts.googleapis.com/css2?family=IBM+Plex+Sans:wght@400;500;600;700&family=IBM+Plex+Mono:wght@400;500&family=Source+Serif+4:wght@500;600;700&display=swap');
+
+:root {
+  --paper: #f4efe6;
+  --ink: #122033;
+  --ink-soft: #4f6073;
+  --accent: #0f5fd7;
+  --accent-soft: #d9e8ff;
+  --line: #c9d1db;
+  --warning: #8a4b0f;
+  --good: #0e6b48;
+}
+
+.gradio-container {
+  font-family: "IBM Plex Sans", sans-serif;
+  background:
+    radial-gradient(circle at top right, rgba(15,95,215,0.08), transparent 24rem),
+    linear-gradient(180deg, #faf7f0 0%, var(--paper) 100%);
+  color: var(--ink);
+}
+
+#hero {
+  border: 1px solid var(--line);
+  background: linear-gradient(135deg, rgba(255,255,255,0.9), rgba(239,245,255,0.92));
+  border-radius: 24px;
+  padding: 1.25rem 1.4rem;
+  box-shadow: 0 20px 40px rgba(18,32,51,0.08);
+}
+
+.eyebrow {
+  font-family: "IBM Plex Mono", monospace;
+  font-size: 0.78rem;
+  letter-spacing: 0.08em;
+  text-transform: uppercase;
+  color: var(--accent);
+}
+
+.hero-title {
+  font-family: "Source Serif 4", serif;
+  font-size: 2.2rem;
+  line-height: 1.05;
+  margin: 0.2rem 0 0.5rem 0;
+}
+
+.hero-copy {
+  color: var(--ink-soft);
+  max-width: 60rem;
+  font-size: 1rem;
+}
+
+.panel-note {
+  border-left: 4px solid var(--accent);
+  background: rgba(15,95,215,0.06);
+  padding: 0.9rem 1rem;
+  border-radius: 12px;
+}
+
+.metric-strip {
+  display: grid;
+  grid-template-columns: repeat(3, minmax(0, 1fr));
+  gap: 0.8rem;
+}
+
+.metric-card {
+  border: 1px solid var(--line);
+  background: rgba(255,255,255,0.75);
+  padding: 0.8rem 0.9rem;
+  border-radius: 16px;
+}
+
+.metric-card strong {
+  display: block;
+  font-size: 1.15rem;
+  margin-top: 0.15rem;
+}
+"""
+
+EXAMPLES = {
+    "BTK binding": {
+        "title": "BTK kinase inhibitor binding assay",
+        "description": "In vitro kinase-domain binding assay for Bruton's tyrosine kinase inhibitor ranking.",
+        "organism": "Homo sapiens",
+        "readout": "binding",
+        "assay_format": "biochemical",
+        "assay_type": "binding",
+        "target_uniprot": "Q06187",
+        "smiles": "\n".join(
+            [
+                "CC1=NC(=O)N(C)C(=O)N1",
+                "CCOc1ccc2nc(N3CCN(C)CC3)n(C)c(=O)c2c1",
+                "CC(=O)Nc1ncc(C#N)c(Nc2ccc(F)c(Cl)c2)n1",
+                "c1ccccc1",
+                "CCO",
+            ]
+        ),
+    },
+    "ALDH1A1 fluorescence": {
+        "title": "ALDH1A1 inhibition assay",
+        "description": "Cell-based fluorescence assay measuring ALDH1A1 inhibition in human cells.",
+        "organism": "Homo sapiens",
+        "readout": "fluorescence",
+        "assay_format": "cell-based",
+        "assay_type": "inhibition",
+        "target_uniprot": "P00352",
+        "smiles": "\n".join(
+            [
+                "CC1=CC(=O)N(C)C(=O)N1",
+                "COC1=CC=C(C=C1)C(=O)O",
+                "CCN(CC)CCOC1=CC=CC=C1",
+                "CCOC1=CC=CC=C1",
+                "CCO",
+            ]
+        ),
+    },
+}
+
+
+def _parse_smiles_text(value: str | None) -> list[str]:
+    if not value:
+        return []
+    lines = [line.strip() for line in value.replace(",", "\n").splitlines()]
+    return [line for line in lines if line]
+
+
+def _read_uploaded_smiles(file_obj: Any) -> list[str]:
+    if file_obj is None:
+        return []
+    path = Path(file_obj.name if hasattr(file_obj, "name") else str(file_obj))
+    suffix = path.suffix.lower()
+    if suffix in {".txt", ".smi", ".smiles"}:
+        return [line.strip() for line in path.read_text().splitlines() if line.strip()]
+    if suffix == ".csv":
+        frame = pd.read_csv(path)
+        for column in ("smiles", "canonical_smiles", "SMILES"):
+            if column in frame.columns:
+                return [str(item).strip() for item in frame[column].tolist() if str(item).strip()]
+        first = frame.columns[0]
+        return [str(item).strip() for item in frame[first].tolist() if str(item).strip()]
+    raise gr.Error("Upload a .csv, .txt, .smi, or .smiles file.")
+
+
+def _collect_smiles(smiles_text: str, upload_file: Any) -> tuple[list[str], str | None]:
+    items = _parse_smiles_text(smiles_text) + _read_uploaded_smiles(upload_file)
+    deduped: list[str] = []
+    seen: set[str] = set()
+    for item in items:
+        if item not in seen:
+            deduped.append(item)
+            seen.add(item)
+    warning = None
+    if len(deduped) > MAX_INPUT_SMILES:
+        warning = f"Input truncated to the first {MAX_INPUT_SMILES} unique SMILES for interactive use."
+        deduped = deduped[:MAX_INPUT_SMILES]
+    return deduped, warning
+
+
+def _load_model():
+    return load_compatibility_model_from_hub(MODEL_REPO_ID)
+
+
+def _build_summary(query_text: str, valid_rows: list[dict[str, Any]], invalid_rows: list[dict[str, Any]], warning: str | None) -> str:
+    best = valid_rows[0] if valid_rows else None
+    chunks = [
+        "### Run Summary",
+        f"- Model repo: `{MODEL_REPO_ID}`",
+        f"- Assay prompt length: `{len(query_text.split())}` tokens-equivalent words",
+        f"- Valid molecules ranked: `{len(valid_rows)}`",
+        f"- Invalid molecules rejected: `{len(invalid_rows)}`",
+    ]
+    if best is not None:
+        chunks.append(f"- Top hit: `{best['canonical_smiles']}` with score `{best['score']:.3f}`")
+    if warning:
+        chunks.append(f"- Warning: {warning}")
+    chunks.append("")
+    chunks.append("Higher scores mean the model ranks the molecule as more compatible with this assay than lower-scored candidates in the same list. Scores are ranking signals, not calibrated probabilities.")
+    return "\n".join(chunks)
+
+
+def _results_to_csv(valid_rows: list[dict[str, Any]], invalid_rows: list[dict[str, Any]]) -> str | None:
+    rows = valid_rows + invalid_rows
+    if not rows:
+        return None
+    handle = tempfile.NamedTemporaryFile("w", suffix=".csv", delete=False, newline="")
+    writer = csv.DictWriter(handle, fieldnames=["rank", "input_smiles", "canonical_smiles", "smiles_hash", "score", "valid", "error"])
+    writer.writeheader()
+    rank = 1
+    for row in valid_rows:
+        writer.writerow(
+            {
+                "rank": rank,
+                "input_smiles": row["input_smiles"],
+                "canonical_smiles": row["canonical_smiles"],
+                "smiles_hash": row["smiles_hash"],
+                "score": row["score"],
+                "valid": True,
+                "error": "",
+            }
+        )
+        rank += 1
+    for row in invalid_rows:
+        writer.writerow(
+            {
+                "rank": "",
+                "input_smiles": row["input_smiles"],
+                "canonical_smiles": "",
+                "smiles_hash": "",
+                "score": "",
+                "valid": False,
+                "error": row.get("error", "invalid_smiles"),
+            }
+        )
+    handle.close()
+    return handle.name
+
+
+def run_ranking(
+    title: str,
+    description: str,
+    organism: str,
+    readout: str,
+    assay_format: str,
+    assay_type: str,
+    target_uniprot: str,
+    smiles_text: str,
+    upload_file: Any,
+    top_k: int,
+):
+    smiles_values, warning = _collect_smiles(smiles_text, upload_file)
+    if not smiles_values:
+        raise gr.Error("Provide at least one SMILES entry by paste or file upload.")
+    query = AssayQuery(
+        title=title or "",
+        description=description or "",
+        organism=organism or "",
+        readout=readout or "",
+        assay_format=assay_format or "",
+        assay_type=assay_type or "",
+        target_uniprot=[token.strip() for token in target_uniprot.split(",") if token.strip()],
+    )
+    assay_text = serialize_assay_query(query)
+    model = _load_model()
+    ranked = rank_compounds(model, assay_text=assay_text, smiles_list=smiles_values, top_k=top_k or None)
+    valid_rows = [row for row in ranked if row["valid"]]
+    invalid_rows = [row for row in ranked if not row["valid"]]
+
+    display_rows = [
+        {
+            "rank": idx + 1,
+            "input_smiles": row["input_smiles"],
+            "canonical_smiles": row["canonical_smiles"],
+            "smiles_hash": row["smiles_hash"],
+            "score": round(float(row["score"]), 4),
+        }
+        for idx, row in enumerate(valid_rows)
+    ]
+    invalid_display = [
+        {"input_smiles": row["input_smiles"], "error": row.get("error", "invalid_smiles")}
+        for row in invalid_rows
+    ]
+    summary = _build_summary(assay_text, valid_rows, invalid_rows, warning)
+    csv_path = _results_to_csv(valid_rows, invalid_rows)
+    return summary, assay_text, pd.DataFrame(display_rows), pd.DataFrame(invalid_display), csv_path
+
+
+def load_example(example_name: str):
+    example = EXAMPLES[example_name]
+    return (
+        example["title"],
+        example["description"],
+        example["organism"],
+        example["readout"],
+        example["assay_format"],
+        example["assay_type"],
+        example["target_uniprot"],
+        example["smiles"],
+    )
+
+
+with gr.Blocks(css=CSS, title="BioAssayAlign Compatibility Explorer") as demo:
+    gr.Markdown(
+        """
+<div id="hero">
+  <div class="eyebrow">BioAssayAlign · scientist-facing ranking demo</div>
+  <div class="hero-title">Rank candidate molecules for a bioassay</div>
+  <div class="hero-copy">
+    Build an assay query from structured fields, paste or upload a candidate molecule list, and get a ranked output from the current BioAssayAlign compatibility model.
+    This app is designed for triage and prioritization, not for direct potency claims.
+  </div>
+</div>
+"""
+    )
+
+    with gr.Row():
+        with gr.Column(scale=5):
+            gr.Markdown(
+                """
+<div class="panel-note">
+Use the structured fields if you have them. Missing fields are allowed, but species, readout, and target metadata usually help.
+</div>
+"""
+            )
+        with gr.Column(scale=4):
+            gr.Markdown(
+                f"""
+<div class="metric-strip">
+  <div class="metric-card"><span>Default model</span><strong>{MODEL_REPO_ID}</strong></div>
+  <div class="metric-card"><span>Expected use</span><strong>ranking, not probability</strong></div>
+  <div class="metric-card"><span>Interactive cap</span><strong>{MAX_INPUT_SMILES} SMILES</strong></div>
+</div>
+"""
+            )
+
+    with gr.Tab("Rank Compounds"):
+        with gr.Row():
+            with gr.Column(scale=6):
+                example_name = gr.Dropdown(choices=list(EXAMPLES.keys()), value="BTK binding", label="Load an example")
+                load_example_btn = gr.Button("Load Example", variant="secondary")
+                assay_title = gr.Textbox(label="Assay title")
+                description = gr.Textbox(label="Description", lines=6, placeholder="Describe the assay in practical lab language.")
+                with gr.Row():
+                    organism = gr.Textbox(label="Organism", placeholder="Homo sapiens")
+                    readout = gr.Textbox(label="Readout", placeholder="binding / fluorescence / luminescence")
+                with gr.Row():
+                    assay_format = gr.Textbox(label="Assay format", placeholder="biochemical / cell-based")
+                    assay_type = gr.Textbox(label="Assay type", placeholder="binding / inhibition / activation")
+                target_uniprot = gr.Textbox(label="Target UniProt IDs", placeholder="Q06187, P52333")
+
+            with gr.Column(scale=5):
+                smiles_text = gr.Textbox(
+                    label="Candidate SMILES",
+                    lines=14,
+                    placeholder="Paste one SMILES per line. CSV upload is optional and will be merged.",
+                )
+                upload_file = gr.File(label="Upload CSV / TXT / SMI", file_count="single", file_types=[".csv", ".txt", ".smi", ".smiles"])
+                top_k = gr.Slider(label="Top-K rows to display", minimum=5, maximum=200, step=5, value=DEFAULT_TOP_K)
+                run_btn = gr.Button("Rank Molecules", variant="primary")
+                clear_btn = gr.ClearButton(value="Clear", components=[assay_title, description, organism, readout, assay_format, assay_type, target_uniprot, smiles_text, upload_file])
+
+        summary = gr.Markdown()
+        with gr.Accordion("Serialized assay text used by the model", open=False):
+            assay_preview = gr.Textbox(lines=12, show_copy_button=True, label="Model-facing assay text")
+        ranked_df = gr.Dataframe(label="Ranked molecules", interactive=False, wrap=True)
+        invalid_df = gr.Dataframe(label="Rejected inputs", interactive=False, wrap=True)
+        download_file = gr.File(label="Download CSV")
+
+        load_example_btn.click(
+            load_example,
+            inputs=[example_name],
+            outputs=[assay_title, description, organism, readout, assay_format, assay_type, target_uniprot, smiles_text],
+        )
+        run_btn.click(
+            run_ranking,
+            inputs=[assay_title, description, organism, readout, assay_format, assay_type, target_uniprot, smiles_text, upload_file, top_k],
+            outputs=[summary, assay_preview, ranked_df, invalid_df, download_file],
+        )
+
+    with gr.Tab("How To Use This"):
+        gr.Markdown(
+            """
+### Recommended workflow
+
+1. Describe the assay in plain scientific language.
+2. Add metadata if you know it: organism, readout, format, assay type, target UniProt.
+3. Paste a candidate list or upload a CSV with a `smiles` column.
+4. Rank the list and inspect the top molecules first.
+
+### What the score means
+
+- The score is a ranking signal.
+- Higher means “more compatible than the other molecules in this submitted list”.
+- It is **not** a calibrated activity probability and it is **not** an IC50 prediction.
+
+### Good input habits
+
+- Prefer parent, neutralized, chemically sensible SMILES.
+- Keep assay descriptions concrete.
+- If the assay is target-defined, add the UniProt ID.
+
+### What this Space is not
+
+- not a generative chemistry tool
+- not a medicinal chemistry oracle
+- not a wet-lab substitute
+"""
+        )
+
+
+if __name__ == "__main__":
+    demo.queue(default_concurrency_limit=4).launch()

examples/btk_candidates.csv

@@ -0,0 +1,6 @@
+smiles
+CC1=NC(=O)N(C)C(=O)N1
+CCOc1ccc2nc(N3CCN(C)CC3)n(C)c(=O)c2c1
+CC(=O)Nc1ncc(C#N)c(Nc2ccc(F)c(Cl)c2)n1
+c1ccccc1
+CCO

requirements.txt

@@ -0,0 +1,8 @@
+gradio>=5.0,<6
+huggingface_hub>=0.30
+numpy<2
+pandas>=2.2
+rdkit-pypi>=2022.9.5
+sentence-transformers>=5.2
+torch>=2.2
+transformers>=4.51

space_runtime.py

@@ -0,0 +1,554 @@
+from __future__ import annotations
+
+import hashlib
+import json
+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 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 import AutoModel, AutoTokenizer
+
+RDLogger.DisableLog("rdApp.*")
+
+DEFAULT_ASSAY_TASK = (
+    "Given a bioassay description and metadata, represent the assay for ranking compatible small molecules."
+)
+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}
+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",
+}
+
+
+@dataclass
+class AssayQuery:
+    title: str = ""
+    description: str = ""
+    organism: str = ""
+    readout: str = ""
+    assay_format: str = ""
+    assay_type: str = ""
+    target_uniprot: list[str] | None = None
+
+
+def smiles_sha256(smiles: str) -> str:
+    return hashlib.sha256(smiles.encode("utf-8")).hexdigest()
+
+
+@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 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 _hash_bucket(value: str, dim: int) -> int:
+    return abs(hash(value)) % max(dim, 1)
+
+
+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 _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
+
+
+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 _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 _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.array([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
+
+
+class SpaceCompatibilityModel:
+    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,
+        molecule_transformer_model_name: str,
+        molecule_transformer_batch_size: int,
+        molecule_transformer_max_length: int,
+        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.molecule_transformer_model_name = molecule_transformer_model_name
+        self.molecule_transformer_batch_size = molecule_transformer_batch_size
+        self.molecule_transformer_max_length = molecule_transformer_max_length
+        self.use_assay_metadata_features = use_assay_metadata_features
+        self.assay_metadata_dim = assay_metadata_dim
+        self._molecule_transformer_tokenizer = None
+        self._molecule_transformer_model = None
+        self._molecule_transformer_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    def _format_assay_query(self, assay_text: str) -> str:
+        return f"Instruct: {self.assay_task_description.strip()}\nQuery: {assay_text.strip()}"
+
+    def _build_assay_feature_array(self, assay_text: str) -> np.ndarray:
+        assay_features = self.assay_encoder.encode(
+            [self._format_assay_query(assay_text)],
+            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 _ensure_molecule_transformer_loaded(self) -> None:
+        if not self.molecule_transformer_model_name or self._molecule_transformer_model is not None:
+            return
+        dtype = torch.float16 if self._molecule_transformer_device.type == "cuda" else torch.float32
+        self._molecule_transformer_tokenizer = AutoTokenizer.from_pretrained(
+            self.molecule_transformer_model_name,
+            trust_remote_code=True,
+        )
+        self._molecule_transformer_model = AutoModel.from_pretrained(
+            self.molecule_transformer_model_name,
+            trust_remote_code=True,
+            torch_dtype=dtype,
+        ).to(self._molecule_transformer_device)
+        self._molecule_transformer_model.eval()
+
+    def _encode_molecule_transformer_batch(self, smiles_values: list[str]) -> np.ndarray | None:
+        if not self.molecule_transformer_model_name:
+            return None
+        self._ensure_molecule_transformer_loaded()
+        assert self._molecule_transformer_model is not None
+        assert self._molecule_transformer_tokenizer is not None
+        outputs: list[np.ndarray] = []
+        batch_size = max(self.molecule_transformer_batch_size, 1)
+        with torch.no_grad():
+            for start in range(0, len(smiles_values), batch_size):
+                batch = smiles_values[start : start + batch_size]
+                encoded = self._molecule_transformer_tokenizer(
+                    batch,
+                    padding=True,
+                    truncation=True,
+                    max_length=self.molecule_transformer_max_length,
+                    return_tensors="pt",
+                )
+                encoded = {key: value.to(self._molecule_transformer_device) for key, value in encoded.items()}
+                hidden = self._molecule_transformer_model(**encoded).last_hidden_state
+                mask = encoded["attention_mask"].unsqueeze(-1)
+                pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
+                outputs.append(pooled.detach().cpu().to(torch.float32).numpy())
+        return np.concatenate(outputs, axis=0).astype(np.float32)
+
+    def build_molecule_feature_matrix(self, smiles_values: list[str]) -> np.ndarray:
+        transformer_matrix = self._encode_molecule_transformer_batch(smiles_values)
+        rows: list[np.ndarray] = []
+        for idx, smiles in enumerate(smiles_values):
+            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=self.fingerprint_bits, use_chirality=self.use_chirality)
+                for radius in self.fingerprint_radii
+            ]
+            if self.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 self.use_rdkit_descriptors and self.descriptor_names:
+                dense = _molecule_descriptor_vector(mol, names=self.descriptor_names)
+                if self.descriptor_mean is not None and self.descriptor_std is not None:
+                    dense = (dense - self.descriptor_mean) / self.descriptor_std
+                output_blocks.append(dense.astype(np.float32))
+            if transformer_matrix is not None:
+                output_blocks.append(np.asarray(transformer_matrix[idx], dtype=np.float32))
+            rows.append(np.concatenate(output_blocks, axis=0).astype(np.float32))
+        return np.stack(rows, axis=0)
+
+
+def _load_sentence_transformer(model_name: str) -> SentenceTransformer:
+    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,
+        "molecule_transformer_model_name": str(cfg.get("molecule_transformer_model_name") or ""),
+        "molecule_transformer_max_length": int(cfg.get("molecule_transformer_max_length", 128) or 128),
+    }
+
+
+def load_compatibility_model(model_dir: str | Path) -> SpaceCompatibilityModel:
+    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(
+            f"Checkpoint mismatch: unexpected={sorted(unexpected)} missing={sorted(missing)}"
+        )
+    return SpaceCompatibilityModel(
+        assay_encoder=encoder,
+        compatibility_head=head,
+        assay_task_description=checkpoint.get("assay_task_description") or cfg.get("assay_task_description", DEFAULT_ASSAY_TASK),
+        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", cfg.get("use_maccs", False))),
+        use_rdkit_descriptors=bool(feature_spec.get("use_rdkit_descriptors", cfg.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,
+        molecule_transformer_model_name=str(feature_spec.get("molecule_transformer_model_name") or cfg.get("molecule_transformer_model_name") or ""),
+        molecule_transformer_batch_size=int(cfg.get("molecule_transformer_batch_size", 128) or 128),
+        molecule_transformer_max_length=int(feature_spec.get("molecule_transformer_max_length") or cfg.get("molecule_transformer_max_length", 128) or 128),
+        use_assay_metadata_features=bool(cfg.get("use_assay_metadata_features", False)),
+        assay_metadata_dim=int(cfg.get("assay_metadata_dim", 0) or 0),
+    )
+
+
+@lru_cache(maxsize=1)
+def load_compatibility_model_from_hub(model_repo_id: str) -> SpaceCompatibilityModel:
+    model_dir = snapshot_download(
+        repo_id=model_repo_id,
+        repo_type="model",
+        allow_patterns=["best_model.pt", "training_metadata.json", "README.md"],
+    )
+    return load_compatibility_model(model_dir)
+
+
+def rank_compounds(
+    model: SpaceCompatibilityModel,
+    *,
+    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