Upload bioassayalign_compatibility.py with huggingface_hub

bioassayalign_compatibility.py

@@ -39,7 +39,9 @@ SECTION_ORDER = [
     "ASSAY_TYPE",
     "TARGET_UNIPROT",
 ]
-ASSAY_SECTION_RE = re.compile(r"\[(ASSAY_TITLE|DESCRIPTION|ORGANISM|READOUT|ASSAY_FORMAT|ASSAY_TYPE|TARGET_UNIPROT)\]\n")
+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",
@@ -106,6 +108,10 @@ def serialize_assay_query(query: AssayQuery) -> str:
     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)
@@ -219,7 +225,9 @@ def _molecule_descriptor_vector(mol, *, names: tuple[str, ...] = DEFAULT_DESCRIP
     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)
+    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 = {
@@ -258,177 +266,304 @@ def _molecule_descriptor_vector(mol, *, names: tuple[str, ...] = DEFAULT_DESCRIP
 
 
 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:
+    def __init__(self, *, assay_dim: int, molecule_dim: int, projection_dim: int, hidden_dim: int, dropout: float) -> 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(
+        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),
         )
-
-    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(
+        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_hidden,
-                molecule_hidden,
-                assay_hidden * molecule_hidden,
-                torch.abs(assay_hidden - molecule_hidden),
+                assay_expand,
+                mol_vec,
+                assay_expand * mol_vec,
+                torch.abs(assay_expand - mol_vec),
             ],
             dim=-1,
         )
-        return self.scorer(interaction).squeeze(-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 CompatibilityModel:
-    def __init__(self, assay_encoder: SentenceTransformer, metadata: dict[str, Any], model_state_dict: 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")
+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.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(model_state_dict)
-        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,
+        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,
-            prompt_name="query",
-            prompt=DEFAULT_ASSAY_TASK,
+            show_progress_bar=False,
+            convert_to_numpy=True,
         )[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)
+        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 load_compatibility_model(model_dir: str | Path, *, device: str | None = None) -> CompatibilityModel:
+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)
-    training_metadata = json.loads((model_path / "training_metadata.json").read_text())
     checkpoint = torch.load(model_path / "best_model.pt", map_location="cpu", weights_only=False)
-    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["model_state_dict"], 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"})
+    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_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(
+        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": input_smiles,
-                    "canonical_smiles": standardized,
-                    "smiles_hash": smiles_sha256(standardized),
+                    "input_smiles": raw_smiles,
+                    "canonical_smiles": canonical,
+                    "smiles_hash": smiles_sha256(canonical),
                     "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
+        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]: