"""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, } # --------------------------------------------------------------------------- # Peptide prior # --------------------------------------------------------------------------- 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." ) # --------------------------------------------------------------------------- # Anchor prior # --------------------------------------------------------------------------- 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)), ) # merge default weights with user overrides 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." ) # --------------------------------------------------------------------------- # Block prior # --------------------------------------------------------------------------- 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." ) # --------------------------------------------------------------------------- # Bundle # --------------------------------------------------------------------------- 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