| """Factories for peptide / anchor / block reference priors. |
| |
| The training script no longer instantiates ``MockPeptidePrior`` / |
| ``MockAnchorPrior`` / ``MockBlockPrior`` directly. Instead it calls |
| :func:`build_reference_priors`, which reads the ``reference_priors`` block |
| from the training YAML and dispatches to the requested backend. |
| |
| Supported backends |
| ------------------ |
| peptide : ``mock`` | ``heuristic`` | ``esm2_delta`` |
| anchor : ``mock`` | ``motif_support_geometry`` |
| block : ``mock`` | ``catalog_scored`` |
| |
| ``strict_no_mock=true`` at the top level forces every backend to be a real |
| prior; a mock backend triggers :class:`ValueError` (surfaced by the trainer as |
| a ``SystemExit``). |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from typing import Any |
|
|
| from staplebridge.data.schemas import BuildingBlock |
| from staplebridge.oracles.anchor_prior import MockAnchorPrior |
| from staplebridge.oracles.base import ( |
| AnchorPriorBase, |
| BlockPriorBase, |
| GeometryOracleBase, |
| PeptidePriorBase, |
| ) |
| from staplebridge.oracles.block_prior import MockBlockPrior |
| from staplebridge.oracles.catalog_block_prior import ( |
| CatalogBlockPrior, |
| CatalogBlockPriorConfig, |
| ) |
| from staplebridge.oracles.esm2_peptide_prior import ESM2DeltaPeptidePrior |
| from staplebridge.oracles.heuristic_peptide_prior import HeuristicPeptidePrior |
| from staplebridge.oracles.motif_anchor_prior import ( |
| MotifSupportAnchorConfig, |
| MotifSupportAnchorPrior, |
| ) |
| from staplebridge.oracles.peptide_prior import MockPeptidePrior |
|
|
|
|
| MOCK_PEPTIDE_BACKENDS = {"mock", "heuristic"} |
| MOCK_ANCHOR_BACKENDS = {"mock"} |
| MOCK_BLOCK_BACKENDS = {"mock"} |
|
|
|
|
| @dataclass |
| class PriorSelection: |
| peptide: PeptidePriorBase |
| anchor: AnchorPriorBase |
| block: BlockPriorBase |
| peptide_backend: str |
| anchor_backend: str |
| block_backend: str |
| strict_no_mock: bool = False |
|
|
| @property |
| def any_mock(self) -> bool: |
| return ( |
| self.peptide_backend in MOCK_PEPTIDE_BACKENDS |
| or self.anchor_backend in MOCK_ANCHOR_BACKENDS |
| or self.block_backend in MOCK_BLOCK_BACKENDS |
| ) |
|
|
| def to_manifest(self) -> dict[str, Any]: |
| return { |
| "peptide_backend": self.peptide_backend, |
| "anchor_backend": self.anchor_backend, |
| "block_backend": self.block_backend, |
| "strict_no_mock": bool(self.strict_no_mock), |
| "peptide_is_mock": self.peptide_backend in MOCK_PEPTIDE_BACKENDS, |
| "anchor_is_mock": self.anchor_backend in MOCK_ANCHOR_BACKENDS, |
| "block_is_mock": self.block_backend in MOCK_BLOCK_BACKENDS, |
| } |
|
|
|
|
| |
| |
| |
| def build_peptide_prior( |
| cfg: dict[str, Any] | None, |
| sequence_encoder: Any = None, |
| device: str = "cpu", |
| ) -> tuple[PeptidePriorBase, str]: |
| cfg = dict(cfg or {}) |
| backend = str(cfg.get("backend", "mock")).lower() |
| if backend == "mock": |
| return MockPeptidePrior(), "mock" |
| if backend == "heuristic": |
| return HeuristicPeptidePrior( |
| target_length=int(cfg.get("target_length", 12)), |
| length_weight=float(cfg.get("length_weight", 1.0)), |
| risky_penalty=float(cfg.get("risky_penalty", 0.3)), |
| hydrophobic_target=float(cfg.get("hydrophobic_target", 0.45)), |
| hydrophobic_weight=float(cfg.get("hydrophobic_weight", 0.4)), |
| ), "heuristic" |
| if backend == "esm2_delta": |
| prior = ESM2DeltaPeptidePrior( |
| model_name_or_path=cfg.get("model_name_or_path"), |
| sequence_encoder=sequence_encoder, |
| cache_path=cfg.get("cache_path"), |
| ncaa_policy=cfg.get("ncaa_policy"), |
| offline=bool(cfg.get("offline", False)), |
| device=str(cfg.get("device", device)), |
| temperature=float(cfg.get("temperature", 1.0)), |
| max_batch_size=int(cfg.get("max_batch_size", 128)), |
| memory_cache_size=int(cfg.get("memory_cache_size", 65536)), |
| strict_runtime=bool(cfg.get("strict_runtime", False)), |
| ) |
| return prior, "esm2_delta" |
| raise ValueError( |
| f"Unknown reference_priors.peptide.backend: {backend!r}. " |
| f"Expected one of: mock, heuristic, esm2_delta." |
| ) |
|
|
|
|
| |
| |
| |
| def build_anchor_prior( |
| cfg: dict[str, Any] | None, |
| geometry_oracle: GeometryOracleBase | None = None, |
| catalog: dict[str, BuildingBlock] | None = None, |
| ) -> tuple[AnchorPriorBase, str]: |
| cfg = dict(cfg or {}) |
| backend = str(cfg.get("backend", "mock")).lower() |
| if backend == "mock": |
| return MockAnchorPrior(), "mock" |
| if backend == "motif_support_geometry": |
| weights = dict(cfg.get("weights") or {}) |
| conf = MotifSupportAnchorConfig( |
| allowed_spacings=tuple(cfg.get("allowed_spacings") or (3, 4)), |
| max_motif_edits=int(cfg.get("max_motif_edits", 2)), |
| avoid_protected=bool(cfg.get("avoid_protected", True)), |
| project_motif_for_geometry=bool(cfg.get("project_motif_for_geometry", False)), |
| ) |
| |
| conf.weights.update({k: float(v) for k, v in weights.items()}) |
| return MotifSupportAnchorPrior( |
| catalog=catalog, |
| geometry_oracle=geometry_oracle, |
| cfg=conf, |
| ), "motif_support_geometry" |
| raise ValueError( |
| f"Unknown reference_priors.anchor.backend: {backend!r}. " |
| f"Expected one of: mock, motif_support_geometry." |
| ) |
|
|
|
|
| |
| |
| |
| def build_block_prior( |
| cfg: dict[str, Any] | None, |
| catalog: dict[str, BuildingBlock] | None = None, |
| geometry_oracle: GeometryOracleBase | None = None, |
| ) -> tuple[BlockPriorBase, str]: |
| cfg = dict(cfg or {}) |
| backend = str(cfg.get("backend", "mock")).lower() |
| if backend == "mock": |
| return MockBlockPrior(), "mock" |
| if backend == "catalog_scored": |
| weights = dict(cfg.get("weights") or {}) |
| conf = CatalogBlockPriorConfig( |
| max_motif_edits=int(cfg.get("max_motif_edits", 2)), |
| no_block_penalty=float(cfg.get("no_block_penalty", 0.3)), |
| no_anchor_penalty=float(cfg.get("no_anchor_penalty", 1.0)), |
| ) |
| conf.weights.update({k: float(v) for k, v in weights.items()}) |
| return CatalogBlockPrior( |
| catalog=catalog, |
| geometry_oracle=geometry_oracle, |
| cfg=conf, |
| ), "catalog_scored" |
| raise ValueError( |
| f"Unknown reference_priors.block.backend: {backend!r}. " |
| f"Expected one of: mock, catalog_scored." |
| ) |
|
|
|
|
| |
| |
| |
| def build_reference_priors( |
| cfg: dict[str, Any] | None, |
| catalog: dict[str, BuildingBlock] | None, |
| geometry_oracle: GeometryOracleBase | None, |
| sequence_encoder: Any = None, |
| device: str = "cpu", |
| ) -> PriorSelection: |
| """Build all three priors from the ``reference_priors`` config section. |
| |
| ``cfg`` can be ``None`` (meaning "legacy: everything Mock"), a dict with |
| ``peptide/anchor/block`` sub-blocks, or a dict where those sub-blocks are |
| missing (they default to ``backend: mock``). |
| |
| ``strict_no_mock=true`` raises :class:`ValueError` if any resolved backend |
| is a mock/heuristic; the trainer surfaces this as :class:`SystemExit`. |
| """ |
| cfg = dict(cfg or {}) |
| strict = bool(cfg.get("strict_no_mock", False)) |
| peptide, pb = build_peptide_prior( |
| cfg.get("peptide"), sequence_encoder=sequence_encoder, device=device |
| ) |
| anchor, ab = build_anchor_prior( |
| cfg.get("anchor"), geometry_oracle=geometry_oracle, catalog=catalog |
| ) |
| block, bb = build_block_prior( |
| cfg.get("block"), catalog=catalog, geometry_oracle=geometry_oracle |
| ) |
| sel = PriorSelection( |
| peptide=peptide, |
| anchor=anchor, |
| block=block, |
| peptide_backend=pb, |
| anchor_backend=ab, |
| block_backend=bb, |
| strict_no_mock=strict, |
| ) |
| if strict and sel.any_mock: |
| offending = [] |
| if pb in MOCK_PEPTIDE_BACKENDS: |
| offending.append(f"peptide={pb}") |
| if ab in MOCK_ANCHOR_BACKENDS: |
| offending.append(f"anchor={ab}") |
| if bb in MOCK_BLOCK_BACKENDS: |
| offending.append(f"block={bb}") |
| raise ValueError( |
| "reference_priors.strict_no_mock=true but the following backends " |
| f"are still mock/heuristic: {', '.join(offending)}. Either switch " |
| "them to real backends (esm2_delta / motif_support_geometry / " |
| "catalog_scored) or set strict_no_mock=false for ablations." |
| ) |
| return sel |
|
|
