Upload Staplebridge files (#1)

bb6d2aa 30 days ago

22.1 kB

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": <dict>, "score": float,
        "property": <std_key>, "mode": <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

	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": <dict>, "score": float,
	"property": <std_key>, "mode": <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