Synthyra/ESMFold2
0498
1# Copyright 2026 Biohub. All rights reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14"""ESMFold2 model configuration."""
15
16from __future__ import annotations
17
18from dataclasses import asdict, dataclass, field
19
20from transformers.configuration_utils import PretrainedConfig
21
22# ---------------------------------------------------------------------------
23# Nested dataclass configs
24# ---------------------------------------------------------------------------
25
26_DEFAULT_ESMC_HF_REPO = "Synthyra/ESMplusplus_6B"
27_DEFAULT_ESMC_ATTN_BACKEND = "flex"
28_ESMC_ID_ALIASES = {
29 "biohub/ESMC-300M": "Synthyra/ESMplusplus_small",
30 "biohub/ESMC-600M": "Synthyra/ESMplusplus_large",
31 "biohub/ESMC-6B": "Synthyra/ESMplusplus_6B",
32}
33
34
35def normalize_esmc_id(esmc_id: str) -> str:
36 if esmc_id in _ESMC_ID_ALIASES:
37 return _ESMC_ID_ALIASES[esmc_id]
38 return esmc_id
39
40
41@dataclass
42class MSAEncoderConfig:
43 """Config for the optional MSA encoder module (Large MSA models only)."""
44
45 enabled: bool = False
46 d_msa: int = 128
47 d_hidden: int = 32
48 n_layers: int = 4
49 n_heads_msa: int = 8
50 msa_head_width: int = 32
51
52
53@dataclass
54class ParcaeConfig:
55 """Release-only config for the parcae diffusion-loop scheduler."""
56
57 enabled: bool = True
58 poisson_mean: float = 3.0
59 min_steps: int = 1
60 max_steps: int | None = 6
61 coda_n_layers: int = 2
62
63
64@dataclass
65class LMEncoderConfig:
66 """Release-only config for the LM-side pair encoder."""
67
68 enabled: bool = True
69 n_layers: int = 4
70 lm_dropout: float = 0.25
71 per_loop_lm_dropout: bool = True
72
73
74@dataclass
75class AtomAttentionConfig:
76 """Config for SWA atom encoder/decoder with 3D RoPE."""
77
78 d_atom: int = 128
79 d_token: int = 768
80 n_blocks: int = 3
81 n_heads: int = 4
82 swa_window_size: int = 128
83 expansion_ratio: int = 2
84 # 3D RoPE config
85 spatial_rope_base_frequency: float = 20.0
86 n_spatial_rope_pairs_per_axis: int = 2
87 n_uid_rope_pairs: int = 10
88 uid_rope_base_frequency: float = 10000.0
89
90
91@dataclass
92class FoldingTrunkConfig:
93 n_layers: int = 24
94 n_heads: int = 8
95 dropout: float = 0.0
96
97
98@dataclass
99class InputsEmbedderConfig:
100 d_inputs: int = 451
101 atom_encoder: AtomAttentionConfig = field(default_factory=AtomAttentionConfig)
102
103 def __post_init__(self):
104 if isinstance(self.atom_encoder, dict):
105 self.atom_encoder = AtomAttentionConfig(**self.atom_encoder)
106
107
108@dataclass
109class DiffusionModuleConfig:
110 """Config for the DiffusionModule."""
111
112 sigma_data: float = 16.0
113 c_atom: int = 128
114 c_token: int = 768
115 c_z: int = 256
116 c_s_inputs: int = 451
117 fourier_dim: int = 256
118 relpos_r_max: int = 32
119 relpos_s_max: int = 2
120 atom_num_blocks: int = 3
121 atom_num_heads: int = 4
122 token_num_blocks: int = 12
123 token_num_heads: int = 16
124 transition_multiplier: int = 2
125
126
127@dataclass
128class DiffusionStructureHeadConfig:
129 """Config for the diffusion-based structure prediction head."""
130
131 diffusion_module: DiffusionModuleConfig = field(
132 default_factory=DiffusionModuleConfig
133 )
134 distogram_bins: int = 128
135
136 # Training noise: sigma ~ sigma_data * exp(mu + sigma * N(0,1))
137 train_noise_log_mean: float = -1.2
138 train_noise_log_std: float = 1.5
139
140 # Sampling defaults (ODE)
141 gamma_0: float = 0.605
142 gamma_min: float = 1.107
143 noise_scale: float = 0.0
144 step_scale: float = 1.0
145
146 # Inference schedule defaults
147 inference_s_max: float = 160.0
148 inference_s_min: float = 4e-4
149 inference_p: float = 8.0
150 inference_num_steps: int = 68
151
152 def __post_init__(self):
153 if isinstance(self.diffusion_module, dict):
154 self.diffusion_module = DiffusionModuleConfig(**self.diffusion_module)
155
156
157@dataclass
158class ConfidenceHeadConfig:
159 enabled: bool = True
160 num_plddt_bins: int = 50
161 num_pde_bins: int = 64
162 num_pae_bins: int = 64
163 min_dist: float = 2.0
164 max_dist: float = 52.0
165 distogram_bins: int = 128
166 folding_trunk: FoldingTrunkConfig = field(
167 default_factory=lambda: FoldingTrunkConfig(n_layers=4)
168 )
169
170 def __post_init__(self):
171 if isinstance(self.folding_trunk, dict):
172 self.folding_trunk = FoldingTrunkConfig(**self.folding_trunk)
173
174
175# ---------------------------------------------------------------------------
176# Top-level config
177# ---------------------------------------------------------------------------
178
179
180class ESMFold2Config(PretrainedConfig):
181 """
182 Configuration for the ESMFold2 structure prediction model.
183
184 Uses SWA atom encoders with 3D RoPE, a diffusion transformer,
185 a folding trunk, and an ESMC 6B PLM backbone.
186
187 Configuration objects inherit from [`PretrainedConfig`] and can be used to control
188 the model outputs. Read the documentation from [`PretrainedConfig`] for more
189 information.
190
191 Args:
192 d_single (`int`, defaults to 384):
193 Dimensionality of single (per-residue) representations.
194 d_pair (`int`, defaults to 256):
195 Dimensionality of pair (residue-residue) representations.
196 n_relative_residx_bins (`int`, defaults to 32):
197 Number of bins for relative residue index encoding.
198 n_relative_chain_bins (`int`, defaults to 2):
199 Number of bins for relative chain encoding.
200 num_loops (`int`, defaults to 10):
201 Number of trunk loops for iterative refinement.
202 num_diffusion_samples (`int`, defaults to 8):
203 Number of parallel structure predictions to generate.
204 lm_dropout (`float`, defaults to 0.0):
205 Dropout probability on LM pair embeddings. When > 0, dropout is
206 applied with ``training=True`` (including at inference) to match
207 the experimental training recipe used by binder design.
208 force_lm_dropout_during_inference (`bool`, defaults to False):
209 When True, apply ``lm_dropout`` even when ``model.eval()`` and
210 ``lm_dropout`` > 0. Binder-design loads set this to True.
211 lm_mask_pct (`float`, defaults to 0.0):
212 Fraction of LM residue tokens randomly replaced with the LM mask
213 token before running the PLM backbone.
214 disable_msa_features (`bool`, defaults to False):
215 When True, zero out MSA-derived ``profile`` and ``deletion_mean``
216 before the inputs embedder (experimental medium/large checkpoints).
217 inputs (`InputsEmbedderConfig`):
218 Configuration for the inputs embedder module.
219 folding_trunk (`FoldingTrunkConfig`):
220 Configuration for the folding trunk.
221 structure_head (`DiffusionStructureHeadConfig`):
222 Configuration for the diffusion-based structure prediction head.
223 confidence_head (`ConfidenceHeadConfig`):
224 Configuration for the confidence prediction head.
225
226 Examples:
227
228 ```python
229 >>> from transformers import ESMFold2Config, ESMFold2ExperimentalModel
230
231 >>> # Initializing an ESMFold2 configuration
232 >>> configuration = ESMFold2Config(type="experimental")
233
234 >>> # Initializing a model (with random weights) from the configuration
235 >>> model = ESMFold2ExperimentalModel(configuration)
236
237 >>> # Accessing the model configuration
238 >>> configuration = model.config
239 ```
240 """
241
242 model_type = "esmfold2"
243 has_no_defaults_at_init = True
244
245 def __init__(self, **kwargs):
246 super().__init__(**kwargs)
247
248 self.type: str = kwargs.get("type", "release")
249 if self.type not in ("release", "experimental"):
250 raise ValueError(
251 f"ESMFold2Config.type must be 'release' or 'experimental', "
252 f"got {self.type!r}"
253 )
254
255 # Top-level scalar fields
256 self.d_single: int = kwargs.get("d_single", 384)
257 self.d_pair: int = kwargs.get("d_pair", 256)
258 self.n_relative_residx_bins: int = kwargs.get("n_relative_residx_bins", 32)
259 self.n_relative_chain_bins: int = kwargs.get("n_relative_chain_bins", 2)
260 self.num_loops: int = kwargs.get("num_loops", 10)
261 self.num_diffusion_samples: int = kwargs.get("num_diffusion_samples", 8)
262 # If True, ``profile`` / ``deletion_mean`` are zeroed before the inputs
263 # embedder.
264 self.disable_msa_features: bool = kwargs.get("disable_msa_features", False)
265 self.lm_dropout: float = kwargs.get("lm_dropout", 0.0)
266 self.force_lm_dropout_during_inference: bool = kwargs.get(
267 "force_lm_dropout_during_inference", False
268 )
269 self.lm_mask_pct: float = kwargs.get("lm_mask_pct", 0.0)
270
271 self.lm_d_model: int = kwargs.get("lm_d_model", 2560)
272 self.lm_num_layers: int = kwargs.get("lm_num_layers", 80)
273 # Backward-compatible field name; values now point to FastPLMs ESM++.
274 raw_esmc_id = (
275 kwargs["esmc_id"] if "esmc_id" in kwargs else _DEFAULT_ESMC_HF_REPO
276 )
277 self.esmc_id: str = normalize_esmc_id(raw_esmc_id)
278 self.esmc_attn_backend: str = (
279 kwargs["esmc_attn_backend"]
280 if "esmc_attn_backend" in kwargs
281 else _DEFAULT_ESMC_ATTN_BACKEND
282 )
283
284 def _init_nested(cls, val):
285 if isinstance(val, cls):
286 return val
287 if isinstance(val, dict):
288 return cls(**val)
289 return cls()
290
291 self.inputs = _init_nested(InputsEmbedderConfig, kwargs.get("inputs"))
292 self.folding_trunk = _init_nested(
293 FoldingTrunkConfig, kwargs.get("folding_trunk")
294 )
295 self.structure_head = _init_nested(
296 DiffusionStructureHeadConfig, kwargs.get("structure_head")
297 )
298 self.confidence_head = _init_nested(
299 ConfidenceHeadConfig, kwargs.get("confidence_head")
300 )
301 self.msa_encoder = _init_nested(MSAEncoderConfig, kwargs.get("msa_encoder"))
302 # Release-only modules — ignored when ``type == "experimental"``.
303 self.parcae = _init_nested(ParcaeConfig, kwargs.get("parcae"))
304 self.lm_encoder = _init_nested(LMEncoderConfig, kwargs.get("lm_encoder"))
305 # If True, MSA encoder output replaces the pair stream; if False, it is added.
306 self.msa_encoder_overwrite: bool = bool(
307 kwargs.get("msa_encoder_overwrite", True)
308 )
309
310 def to_dict(self):
311 output = super().to_dict()
312 output["inputs"] = asdict(self.inputs)
313 output["folding_trunk"] = asdict(self.folding_trunk)
314 output["structure_head"] = asdict(self.structure_head)
315 output["confidence_head"] = asdict(self.confidence_head)
316 output["msa_encoder"] = asdict(self.msa_encoder)
317 output["parcae"] = asdict(self.parcae)
318 output["lm_encoder"] = asdict(self.lm_encoder)
319 return output
320
321
322__all__ = [
323 "ESMFold2Config",
324 "MSAEncoderConfig",
325 "ParcaeConfig",
326 "LMEncoderConfig",
327 "normalize_esmc_id",
328]
329