File size: 9,155 Bytes
bb6d2aa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 | """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
|