File size: 2,525 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 | from __future__ import annotations
from abc import ABC, abstractmethod
from typing import Any
from staplebridge.chemistry.state import StapleState
from staplebridge.data.schemas import BuildingBlock
class PeptidePriorBase(ABC):
@abstractmethod
def score_transition(self, old_state: StapleState, new_state: StapleState, context: dict[str, Any] | None = None) -> float:
pass
def batch_score_transitions(
self,
old_state: StapleState,
new_states: list[StapleState],
context: dict[str, Any] | None = None,
) -> list[float]:
"""Score N candidate transitions from ``old_state`` at once.
Default implementation just loops ``score_transition``; heavyweight
priors (ESM2) override this so they can share one model forward across
all candidates. Non-sequence-changing candidates are expected to
return exactly 0.0.
"""
return [
self.score_transition(old_state, new, context) for new in new_states
]
def prewarm_requests(
self, pairs: list[tuple[StapleState, list[StapleState]]]
) -> None:
"""Prefetch model outputs for many (z, candidates) pairs at once.
Default is a noop — heavy priors (ESM2) override this to run one
batched model forward covering every request across all pairs, so a
subsequent per-pair ``batch_score_transitions`` call becomes a pure
cache-lookup.
"""
del pairs
class AnchorPriorBase(ABC):
@abstractmethod
def score_anchor(self, sequence: list[str], anchor_pair: tuple[int, int] | None, context: dict[str, Any] | None = None) -> float:
pass
class BlockPriorBase(ABC):
@abstractmethod
def score_block(
self,
sequence: list[str],
anchor_pair: tuple[int, int] | None,
block: BuildingBlock | None,
context: dict[str, Any] | None = None,
) -> float:
pass
class GeometryOracleBase(ABC):
@abstractmethod
def ctype(
self,
sequence: list[str],
anchor_pair: tuple[int, int] | None,
block: BuildingBlock | None,
*,
peptide_ca: list[tuple[float, float, float]] | None = None,
) -> bool:
pass
@abstractmethod
def cgeom(
self,
sequence: list[str],
anchor_pair: tuple[int, int] | None,
block: BuildingBlock | None,
*,
peptide_ca: list[tuple[float, float, float]] | None = None,
) -> float:
pass
|