from __future__ import annotations import math import os import sys import types from pathlib import Path from typing import Any, Sequence import torch PROPERTY_ALIASES = { "hemolysis": "hemolysis", "nf": "nf", "non_fouling": "nf", "non-fouling": "nf", "solubility": "solubility", "toxicity": "toxicity", "permeability": "permeability_penetrance", "permeability_penetrance": "permeability_penetrance", "pampa": "permeability_pampa", "permeability_pampa": "permeability_pampa", "caco2": "permeability_caco2", "permeability_caco2": "permeability_caco2", "half_life": "halflife", "halflife": "halflife", "binding_affinity": "binding_affinity", } def _to_finite_float(x: Any) -> float | None: try: v = float(x) except (TypeError, ValueError): return None if math.isnan(v) or math.isinf(v): return None return v def extract_numeric_score(raw: Any, property_key: str) -> float: """Pull the canonical numeric score for ``property_key`` out of ``raw``. For ``binding_affinity`` we prefer ``raw["affinity"]`` then ``raw["score"]``. For other properties we prefer ``raw["score"]``. Falls back to walking other plausible keys if a top-level dict is given. Raises ``ValueError`` if no finite float can be produced. When the priority keys do not yield a value but the fallback dict-walk does, a ``warnings.warn`` is emitted so silent mis-extraction is surfaced. """ import warnings as _warnings pkey = (property_key or "").lower() candidates: list[Any] = [] if isinstance(raw, dict): if pkey == "binding_affinity": candidates += [raw.get("affinity"), raw.get("score")] else: candidates += [raw.get("score"), raw.get(pkey)] candidates += [raw.get("value"), raw.get("prediction"), raw.get("mean")] elif isinstance(raw, (list, tuple)) and raw: candidates.append(raw[0]) else: candidates.append(raw) for c in candidates: v = _to_finite_float(c) if v is not None: return v if isinstance(raw, dict): finite_pairs = [(k, _to_finite_float(v)) for k, v in raw.items()] finite_pairs = [(k, v) for k, v in finite_pairs if v is not None] if finite_pairs: if len(finite_pairs) > 1: _warnings.warn( f"PeptiVerse score extraction for property={property_key!r}: " f"no priority key matched; falling back to dict walk over " f"{[k for k, _ in finite_pairs]} -> picked first finite value " f"{finite_pairs[0][1]} from key {finite_pairs[0][0]!r}. " "This may be wrong — verify the predictor output schema.", stacklevel=2, ) return finite_pairs[0][1] raise ValueError( f"Cannot extract finite numeric score for property={property_key!r} from raw={raw!r}" ) class PeptiVerseUnavailable(RuntimeError): pass class PeptiVerseWrapper: """Local-import wrapper around ChatterjeeLab/PeptiVerse `inference.py`. The HF repo is expected to be cloned to disk at ``peptiverse_root``; we insert that path into ``sys.path`` and import ``inference.PeptiVersePredictor`` directly. No remote HTTP API is used. Parameters ---------- peptiverse_root : path to the local PeptiVerse repo clone. manifest_path : best_models.txt path (defaults to ``peptiverse_root/best_models.txt``). classifier_weight_root : weight root (defaults to ``peptiverse_root``). device : ``"cpu"`` or ``"cuda"``. strict : if True, raise on any init error. uncertainty : default ``uncertainty`` flag passed to the underlying predictor unless overridden per-call. cache_enabled : memoize identical (property, mode, input, target, uncertainty) calls to avoid re-running the underlying torch models. """ def __init__( self, peptiverse_root: str | Path, manifest_path: str | Path | None = None, classifier_weight_root: str | Path | None = None, device: str | None = None, strict: bool = False, uncertainty: bool = False, cache_enabled: bool = True, hf_cache_dir: str | Path | None = None, esm_model_name_or_path: str | Path | None = None, peptideclm_model_name_or_path: str | Path | None = None, chemberta_model_name_or_path: str | Path | None = None, offline: bool = False, batch_size: int = 32, ) -> None: self.root = Path(peptiverse_root) if peptiverse_root else None self.manifest_path = Path(manifest_path) if manifest_path else None self.classifier_weight_root = ( Path(classifier_weight_root) if classifier_weight_root else None ) self.device = device self.strict = strict self.uncertainty = bool(uncertainty) self.cache_enabled = bool(cache_enabled) self.hf_cache_dir = Path(hf_cache_dir) if hf_cache_dir else None self.esm_model_name_or_path = ( Path(esm_model_name_or_path) if esm_model_name_or_path else None ) self.peptideclm_model_name_or_path = ( Path(peptideclm_model_name_or_path) if peptideclm_model_name_or_path else None ) self.chemberta_model_name_or_path = ( Path(chemberta_model_name_or_path) if chemberta_model_name_or_path else None ) self.offline = bool(offline) self.predictor = None self.available = False self.init_error: str | None = None self._cache: dict[tuple, dict[str, Any]] = {} self.cache_hits = 0 self.cache_misses = 0 self.batch_size = max(1, int(batch_size)) self.batch_chunks = 0 self.batch_oom_retries = 0 self.batch_oom_giveups = 0 self._init_predictor() @staticmethod def _snapshot_from_cache(cache_root: Path, model_id: str) -> Path | None: """Resolve a cached Hugging Face repo to a concrete local snapshot.""" repo = cache_root / "hub" / f"models--{model_id.replace('/', '--')}" snapshots = repo / "snapshots" ref = repo / "refs" / "main" if ref.is_file(): revision = ref.read_text(encoding="utf-8").strip() candidate = snapshots / revision if candidate.is_dir(): return candidate.resolve() candidates = sorted(path for path in snapshots.glob("*") if path.is_dir()) if len(candidates) == 1: return candidates[0].resolve() return None def _local_hf_models(self) -> dict[str, Path]: configured = { "esm_name": self.esm_model_name_or_path, "clm_name": self.peptideclm_model_name_or_path, "chemberta_name": self.chemberta_model_name_or_path, } model_ids = { "esm_name": "facebook/esm2_t33_650M_UR50D", "clm_name": "aaronfeller/PeptideCLM-23M-all", "chemberta_name": "DeepChem/ChemBERTa-77M-MLM", } resolved: dict[str, Path] = {} for argument, configured_path in configured.items(): path = configured_path if path is None and self.hf_cache_dir is not None: path = self._snapshot_from_cache( self.hf_cache_dir, model_ids[argument] ) if path is not None: path = path.expanduser().resolve() if not path.is_dir(): raise FileNotFoundError( f"local PeptiVerse HF model directory missing: {path}" ) resolved[argument] = path return resolved def _init_predictor(self) -> None: if self.hf_cache_dir is not None: cache_root = self.hf_cache_dir.expanduser().resolve() # Deliberately override inherited values. A stale user-level HF # cache was the source of tokenizer lookup failures in offline # runs; this wrapper's configured local model root is authoritative. os.environ["HF_HOME"] = str(cache_root) os.environ["HUGGINGFACE_HUB_CACHE"] = str(cache_root / "hub") os.environ["TRANSFORMERS_CACHE"] = str(cache_root / "hub") if self.offline: os.environ["HF_HUB_OFFLINE"] = "1" os.environ["TRANSFORMERS_OFFLINE"] = "1" if self.root is None or not self.root.exists(): self.init_error = f"peptiverse_root does not exist: {self.root}" if self.strict: raise PeptiVerseUnavailable(self.init_error) return if not (self.root / "inference.py").exists(): self.init_error = f"inference.py not found under {self.root}" if self.strict: raise PeptiVerseUnavailable(self.init_error) return sys.path.insert(0, str(self.root)) try: # PeptiVerse imports ``tokenizer.my_tokenizers`` from its own # source tree. Some environments also install an unrelated # top-level ``tokenizer`` package; once that regular package is in # sys.modules it shadows PeptiVerse's namespace directory even # though the PeptiVerse root is first on sys.path. Bind the local # namespace explicitly before importing inference. local_tokenizer = self.root / "tokenizer" if (local_tokenizer / "my_tokenizers.py").is_file(): loaded = sys.modules.get("tokenizer") loaded_paths = [ str(Path(value).resolve()) for value in (getattr(loaded, "__path__", None) or []) ] if str(local_tokenizer.resolve()) not in loaded_paths: for name in list(sys.modules): if name == "tokenizer" or name.startswith("tokenizer."): del sys.modules[name] namespace = types.ModuleType("tokenizer") namespace.__path__ = [str(local_tokenizer)] # type: ignore[attr-defined] namespace.__package__ = "tokenizer" sys.modules["tokenizer"] = namespace from inference import PeptiVersePredictor # type: ignore # Several PeptiVerse MAPIE calibration artifacts were serialized # by a training script where this estimator lived in __main__. # Re-export the identical local inference.py class there before # joblib loads those artifacts. Point predictions do not use MAPIE # when uncertainty=False, but resolving the class avoids noisy # load failures and preserves optional uncertainty availability. try: from inference import PassthroughRegressor # type: ignore main_module = sys.modules.get("__main__") if main_module is not None and not hasattr( main_module, "PassthroughRegressor" ): setattr( main_module, "PassthroughRegressor", PassthroughRegressor, ) except ImportError: pass manifest = self.manifest_path or (self.root / "best_models.txt") weight_root = self.classifier_weight_root or self.root kwargs: dict[str, Any] = { "manifest_path": str(manifest), "classifier_weight_root": str(weight_root), } if self.device is not None: kwargs["device"] = self.device for argument, model_path in self._local_hf_models().items(): kwargs[argument] = str(model_path) self.predictor = PeptiVersePredictor(**kwargs) self.available = True except Exception as exc: # noqa: BLE001 - surface any import/init failure self.init_error = f"{type(exc).__name__}: {exc}" self.available = False if self.strict: raise PeptiVerseUnavailable(self.init_error) from exc @staticmethod def normalize_property_name(key: str) -> str: return PROPERTY_ALIASES.get(key.lower(), key) @staticmethod def extract_numeric_score(raw: Any, property_key: str) -> float: return extract_numeric_score(raw, property_key) def _ensure_available(self) -> None: if not self.available or self.predictor is None: raise PeptiVerseUnavailable( self.init_error or "PeptiVerse predictor is not available" ) def clear_cache(self) -> None: self._cache.clear() self.cache_hits = 0 self.cache_misses = 0 def _cache_key( self, property_key: str, mode: str, input_str: str, target_seq: str, uncertainty: bool ) -> tuple: return (property_key, mode, input_str, target_seq, bool(uncertainty)) def predict_property( self, property_key: str, input_str: str, mode: str = "wt", uncertainty: bool | None = None, ) -> dict[str, Any]: """Run a single non-binding property head. Returns ``{"backend_used": "peptiverse", "raw": , "score": float, "property": , "mode": , "cached": bool}``. """ self._ensure_available() std_key = self.normalize_property_name(property_key) unc = self.uncertainty if uncertainty is None else bool(uncertainty) key = self._cache_key(std_key, mode, input_str, "", unc) if self.cache_enabled and key in self._cache: self.cache_hits += 1 cached = dict(self._cache[key]) cached["cached"] = True return cached raw = self.predictor.predict_property( # type: ignore[union-attr] std_key, mode, input_str, unc ) score = extract_numeric_score(raw, std_key) out = { "backend_used": "peptiverse", "property": std_key, "mode": mode, "raw": raw, "score": float(score), "cached": False, } if self.cache_enabled: self._cache[key] = {k: v for k, v in out.items() if k != "cached"} self.cache_misses += 1 return out def predict_property_batch( self, property_key: str, input_strs: Sequence[str], mode: str = "wt", uncertainty: bool | None = None, batch_size: int | None = None, ) -> list[dict[str, Any]]: """Run one non-binding property head over many inputs. Same per-item contract as :meth:`predict_property`, returned in input order. Duplicates and cache hits are computed once. The speedup comes from batching the *embedding* forward pass; the heads still see one input at a time, so each score is produced by exactly the code path :meth:`predict_property` would have used. Batching an embedder changes the padded sequence length, which perturbs the pooled embedding at ~1e-7. That is why only the embeddings are batched here and the results are cached under the same keys as the scalar path: a later scalar call returns the batch-computed value, so a run cannot mix the two conventions for the same input. Verify numeric agreement for any new property head before relying on it. On CUDA OOM the offending chunk is retried at half the batch size, down to scalar, so a long tail of large molecules degrades rather than fails. """ self._ensure_available() std_key = self.normalize_property_name(property_key) if std_key == "binding_affinity": raise ValueError("use predict_binding_affinity for binding affinity") unc = self.uncertainty if uncertainty is None else bool(uncertainty) inputs = list(input_strs) if not inputs: return [] # Deduplicate while preserving first-seen order; only uncached, unique # inputs reach the backend. pending: list[str] = [] seen: set[str] = set() for value in inputs: key = self._cache_key(std_key, mode, value, "", unc) if self.cache_enabled and key in self._cache: continue if value in seen: continue seen.add(value) pending.append(value) if pending: self._embed_batch(std_key, mode, pending, unc, batch_size) # Every input now either was cached or has a warm embedding, so this # loop is the unmodified scalar path and defines the returned values. return [ self.predict_property(std_key, value, mode=mode, uncertainty=unc) for value in inputs ] def _embed_batch( self, std_key: str, mode: str, pending: Sequence[str], uncertainty: bool, batch_size: int | None, ) -> None: """Warm the underlying embedder's cache for ``pending`` in batches.""" embedder, pooled_kind = self._embedder_for(std_key, mode) if embedder is None: return size = int(batch_size or self.batch_size) index = 0 items = list(pending) while index < len(items): chunk = items[index : index + size] try: self._embed_chunk(embedder, pooled_kind, chunk) except torch.cuda.OutOfMemoryError: torch.cuda.empty_cache() if len(chunk) == 1: # Nothing left to split; let the scalar path surface it. self.batch_oom_giveups += 1 index += 1 continue size = max(1, len(chunk) // 2) self.batch_oom_retries += 1 continue index += len(chunk) self.batch_chunks += 1 def _embedder_for(self, std_key: str, mode: str) -> tuple[Any, str]: """The embedder and pooling kind the head for ``std_key`` will request.""" predictor = self.predictor meta = getattr(predictor, "meta", {}).get((std_key, mode)) if meta is None: return None, "" try: embedder = predictor._get_embedder(meta["emb_tag"]) # type: ignore[union-attr] except (KeyError, ValueError): return None, "" # torch_ckpt heads consume unpooled token embeddings; everything else # consumes the masked mean pool. Mirrors PeptiVersePredictor._get_features. return embedder, ("unpooled" if meta.get("kind") == "torch_ckpt" else "pooled") @staticmethod def _embed_chunk(embedder: Any, pooled_kind: str, chunk: Sequence[str]) -> None: """Forward ``chunk`` once and store per-item results in the embedder cache. Writes into the embedder's own ``_cache_pooled``/``_cache_unpooled`` so the scalar ``pooled``/``unpooled`` calls that follow find them. Items already cached are skipped. """ cache_attr = "_cache_pooled" if pooled_kind == "pooled" else "_cache_unpooled" cache = getattr(embedder, cache_attr, None) if cache is None: return stripped = [str(value).strip() for value in chunk] fresh = [value for value in dict.fromkeys(stripped) if value not in cache] if not fresh: return tokenized = embedder._tokenize(fresh) with torch.no_grad(): hidden = embedder.model( input_ids=tokenized["input_ids"], attention_mask=tokenized["attention_mask"], ).last_hidden_state valid = embedder._valid_mask(tokenized["input_ids"], tokenized["attention_mask"]) if pooled_kind == "pooled": weights = valid.unsqueeze(-1).float() pooled = (hidden * weights).sum(dim=1) / weights.sum(dim=1).clamp(min=1e-9) for position, value in enumerate(fresh): cache[value] = pooled[position : position + 1] return for position, value in enumerate(fresh): row = valid[position] features = hidden[position : position + 1, row, :] mask = torch.ones( (1, features.shape[1]), dtype=torch.bool, device=features.device ) cache[value] = (features, mask) def predict_binding_affinity( self, binder_str: str, target_seq: str, mode: str = "wt", uncertainty: bool | None = None, ) -> dict[str, Any]: """Run the binding-affinity head against a specific protein target.""" self._ensure_available() if not target_seq: raise ValueError("predict_binding_affinity requires target_seq") unc = self.uncertainty if uncertainty is None else bool(uncertainty) key = self._cache_key("binding_affinity", mode, binder_str, target_seq, unc) if self.cache_enabled and key in self._cache: self.cache_hits += 1 cached = dict(self._cache[key]) cached["cached"] = True return cached raw = self.predictor.predict_binding_affinity( # type: ignore[union-attr] mode, target_seq, binder_str, unc ) score = extract_numeric_score(raw, "binding_affinity") out = { "backend_used": "peptiverse", "property": "binding_affinity", "mode": mode, "raw": raw, "score": float(score), "cached": False, } if self.cache_enabled: self._cache[key] = {k: v for k, v in out.items() if k != "cached"} self.cache_misses += 1 return out