Transformers documentation

ESMFold2

Transformers

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.2).

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

该模型于 2026-08-19 贡献至 Hugging Face Transformers。

ESMFold2

概述

ESMFold2 是一款全原子蛋白质结构预测模型。它直接基于氨基酸序列预测三维坐标以及逐残基置信度（pLDDT、PAE、PDE），并以 ESMC 蛋白质语言模型作为其主干网络。该架构结合了采用三维旋转位置嵌入的滑动窗口原子编码器、迭代应用的成对折叠主干、基于扩散的结构头以及置信度头。

模型检查点可在 Hugging Face Hub 上的 biohub/ESMFold2 获取。

使用示例

import torch

from transformers import EsmFold2Model

# The ESMC backbone is bundled in the checkpoint and loaded with the model.
# bf16 is the recommended inference precision.
model = EsmFold2Model.from_pretrained("biohub/ESMFold2", dtype=torch.bfloat16, device_map="auto")

pdb_string = model.infer_protein_as_pdb("MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQ")
print(pdb_string)

infer_protein 会返回原始输出（原子坐标、距离图 logits 和置信度指标），以 EsmFold2Output 的形式提供，供你在需要这些内容而非 PDB 字符串时使用。如果对同一序列多次运行，可能会得到略有不同的预测结果。如果你希望获得完全可复现的结构，请手动设置随机种子。

ESMFold2 每次折叠会生成 config.structure_head.num_diffusion_samples 个结构。infer_protein_as_pdb 会渲染排名最高（pTM 最高）的一个结构；如需选择特定样本，请传入 sample_idx。PDB 文件的 b-factor 列中包含每个残基的 pLDDT，其取值范围与 plddt 输出相同，均为 0-1。

forward 与 fold

结构预测包含两个部分。EsmFold2Model.forward 是第一个部分：它对特征化后的输入执行折叠主干网络，并返回精炼后的成对表示以及距离图，其类型为 EsmFold2TrunkOutput。它不会生成 3D 坐标——ESMFold2 通过迭代去噪获得这些坐标，而该采样循环（包括噪声调度、Kabsch 对齐以及 ODE/SDE 更新）与置信度头调用一起位于 EsmFold2FoldingMixin 中：

方法	用于
`infer_protein_as_pdb(sequence)`	一个 PDB 字符串，直接来自氨基酸序列
`infer_protein(sequence)`	原始的 EsmFold2Output
`fold(**features)`	预特征化输入（`infer_protein` 所调用的）
`forward(**features)`	仅主干部分——一个直方图和一个成对表示，无需采样

调用 fold 或 infer_protein 以获取实际结构。如果只需要 distogram，或者希望自行控制扩散采样器，请使用 forward：fold 会调用一次 forward，然后将其输出传递给 EsmFold2DiffusionModule，后者的 forward 即为单步去噪过程。

使用融合内核实现更快的推理

折叠主干的主要开销在于三角形乘法更新。向 from_pretrained() 传入 use_kernels=True 后，该操作会被替换为通过 kernels 库从 Hub 加载的融合 Triton 内核，且预测结果保持不变。此功能仅适用于推理，且仅限 CUDA；在 CPU 上或未安装该内核时，模型会透明地回退到纯 PyTorch 实现。请确保在内核化发生时模型位于 CUDA 设备上（例如通过 device_map 指定）。

import torch

from transformers import EsmFold2Model

model = EsmFold2Model.from_pretrained(
    "biohub/ESMFold2", dtype=torch.bfloat16, device_map="cuda", use_kernels=True
)

pdb_string = model.infer_protein_as_pdb("MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQ")

EsmFold2Config

class transformers.EsmFold2Config

< source >

参数

hidden_size (int，可选，默认为 384) — 单表示的宽度。
pairwise_hidden_size (int，可选，默认为 256) — 成对表示的宽度。
single_inputs_size (int，可选，默认值为 451) — 输入到主干网络和扩散条件中的拼接单输入特征的宽度： atom_encoder.output_dim 加上两个残基类型 one-hot 编码以及 profile 标量。
pair_transition_intermediate_size (int，可选，默认值为 1024) — 配对流转换的 SwiGLU 宽度。
sliding_window (int，可选，默认值为 128) — 原子堆栈注意力的滑动窗口大小（以 token 索引距离表示），指窗口的总宽度，而非半径。
chunk_size (int，可选，默认为 64) — 用于内存占用较高的 pair-/MSA-stream 操作的块大小。None 表示禁用分块。
num_relative_residx_bins (int，可选，默认为 32) — 相对位置编码中相对残基索引的 bin 数量。
num_relative_chain_bins (int，可选，默认为 2) — 相对位置编码中相对链索引的 bin 数量。
num_loops (int，可选，默认为 10) — 主干细化循环的次数。
num_res_types (int，可选，默认为 33) — 残基类型的数量。
max_atomic_number (int，可选，默认为 128) — 原子特征中元素 one-hot 编码的大小。
char_vocab_size (int, 可选，默认为 64) — 编码原子名称所用的字符词表大小。
max_chars (int，可选，默认为 4) — 每个编码原子名称的字符数。
max_atoms_per_token (int，可选，默认为 23) — 每个 token 的最大原子数。
atom_feature_dim (int，可选) — 原子特征宽度：包括 xyz、电荷和掩码，以及元素和原子名字符的 one-hot 编码。如果未设置，则根据 max_atomic_number、char_vocab_size 和 max_chars 推导得出。
folding_trunk_num_hidden_layers (int，可选，默认为 24) — 折叠主干中成对更新块的数量。
parcae_num_coda_layers (int，可选，默认为 2) — parcae coda 中成对更新块的数量。
atom_encoder (EsmFold2AtomEncoderConfig，可选) — 输入嵌入器 SWA 原子编码器的配置。
structure_head (EsmFold2StructureHeadConfig，可选) — 扩散结构预测头的配置（包含 diffusion_module 子配置）。
confidence_head (EsmFold2ConfidenceHeadConfig，可选) — 置信度头的配置。
msa_encoder (EsmFold2MsaEncoderConfig，可选) — MSA 编码器的配置。
lm_encoder (EsmFold2LmEncoderConfig，可选) — 语言模型隐藏状态编码器的配置。
esmc_config (EsmcConfig，可选) — 用于捆绑的 ESMC 主干网络的配置。

这是用于存储 Esmfold2Model 配置的类。它根据指定的参数实例化 Esmfold2 模型，并定义模型架构。使用默认值实例化配置，将得到与 biohub/ESMFold2 类似的配置。

配置对象继承自 PreTrainedConfig，可用于控制模型输出。有关更多信息，请参阅 PreTrainedConfig 的文档。

validate_architecture

< source >

( )

检查跨越子配置的宽度关系；每个子配置检查其自身的宽度关系。

EsmFold2PreTrainedModel

class transformers.EsmFold2PreTrainedModel

< source >

( config: PreTrainedConfig*inputs**kwargs )

参数

config (PreTrainedConfig) — 模型配置类，包含模型的所有参数。使用配置文件进行初始化时，不会加载与模型关联的权重，仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

该模型继承自 PreTrainedModel。请查阅父类文档，了解该库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头等）。

该模型也是 PyTorch torch.nn.Module 的子类。请将其作为常规的 PyTorch Module 使用，有关通用用法和行为的所有事项，请参阅 PyTorch 文档。

EsmFold2Model

class transformers.EsmFold2Model

< source >

( config: EsmFold2Config )

参数

config (EsmFold2Config) — 模型的配置类，包含模型的所有参数。通过配置文件进行初始化时，不会加载与模型关联的权重，仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

ESMFold2 是一款全原子蛋白质结构预测器，内置了 ESMC 蛋白质语言模型主干。这是标准发布的 ESMFold2 架构，其主干是一个线性循环堆叠（内部称为“parcae”）。

该模型继承自 PreTrainedModel。请查阅父类文档，了解该库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头等）。

该模型也是 PyTorch torch.nn.Module 的子类。请将其作为常规的 PyTorch Module 使用，有关通用用法和行为的所有事项，请参阅 PyTorch 文档。

forward

< source >

( token_index: Tensorresidue_index: Tensorasym_id: Tensorsym_id: Tensorentity_id: Tensormol_type: Tensorres_type: Tensortoken_bonds: Tensorattention_mask: Tensoratom_inputs: EsmFold2AtomInputsdeletion_mean: typing.Optional[torch.Tensor] = Nonemsa: typing.Optional[torch.Tensor] = Nonehas_deletion: typing.Optional[torch.Tensor] = Nonedeletion_value: typing.Optional[torch.Tensor] = Nonemsa_attention_mask: typing.Optional[torch.Tensor] = Noneinput_ids: typing.Optional[torch.Tensor] = Nonelm_hidden_states: typing.Optional[torch.Tensor] = Nonenum_loops: int | None = None**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> ) → [~models.esmfold2.modeling_esmfold2.EsmFold2TrunkOutput] or tuple(torch.Tensor)

参数

token_index（形状为 (batch_size, num_tokens) 的 torch.Tensor）— 完整序列中每个 token 的位置索引；用于相对位置编码。
residue_index（形状为 (batch_size, num_tokens) 的 torch.Tensor）— 每条链内的残基索引；用于相对位置编码。
asym_id（形状为 (batch_size, num_tokens) 的 torch.Tensor）— 每个 token 的非对称单元（链）ID。
sym_id（形状为 (batch_size, num_tokens) 的 torch.Tensor）— 用于区分同源多聚体中相同链的对称拷贝 ID。
entity_id（形状为 (batch_size, num_tokens) 的 torch.Tensor）— 实体 ID，用于将属于同一分子实体的 token 进行分组。
mol_type（形状为 (batch_size, num_tokens) 的 torch.Tensor）— 每个 token 的分子类型编码：0 = 蛋白质，1 = DNA，2 = RNA， 3 = 非聚合物（配体）。
res_type（形状为 (batch_size, num_tokens) 的 torch.Tensor）— 每个 token 的残基类型（氨基酸种类）索引。
token_bonds（形状为 (batch_size, num_tokens, num_tokens, 1) 的 torch.Tensor）— 成对 token 间的共价键特征。
attention_mask（形状为 (batch_size, num_tokens) 的 torch.Tensor）— 用于标记有效 token（1）与填充（0）的掩码。输入必须采用右侧填充。
atom_inputs (~models.esmfold2.modeling_esmfold2.EsmFold2AtomInputs) — 打包的原始参考构象原子张量，包括：位置、原子序数、形式电荷、编码后的原子名称、每个原子的基团 ID、原子填充掩码以及原子到 token 的映射。
deletion_mean（形状为 (batch_size, num_tokens) 的 torch.Tensor，可选） — 每列的平均 MSA 删除计数。默认值为零（无 MSA）。
msa (torch.Tensor，形状为 (batch_size, msa_depth, num_tokens)，可选) — MSA 残基类型 token（第 0 行为查询序列）。默认使用单序列 MSA。
has_deletion (torch.Tensor，形状为 (batch_size, msa_depth, num_tokens)，可选) — 布尔标志，用于标记 MSA 中前面存在缺失（deletion）的位置。
deletion_value（形状为 (batch_size, msa_depth, num_tokens) 的 torch.Tensor，可选） — 每个位置的 MSA 删除计数。
msa_attention_mask（torch.Tensor，形状为 (batch_size, msa_depth, num_tokens)，可选） — 用于 MSA 行/列的有效性掩码。
input_ids (torch.Tensor，形状为 (batch_size, num_tokens)，可选) — 序列的 ESMC 词表 token id。当未直接传入 lm_hidden_states 时，将其输入到内置的 ESMC 主干网络以生成该值；若已提供 lm_hidden_states，则忽略此参数。
lm_hidden_states (torch.Tensor，形状为 (batch_size, num_tokens, num_esmc_layers + 1, esmc_hidden_size)，可选) — 预计算的 ESMC 主干隐藏状态，包含每个主干层以及嵌入层各一个。提供该参数时，将不会运行主干网络，且 input_ids 将被忽略。
num_loops (int, 可选) — 主干细化循环的次数。默认值为 config.num_loops。

返回值

[~models.esmfold2.modeling_esmfold2.EsmFold2TrunkOutput] 或 tuple(torch.Tensor)

一个 [~models.esmfold2.modeling_esmfold2.EsmFold2TrunkOutput] 或 torch.Tensor 元组（如果传入了 return_dict=False 或 config.return_dict=False），其包含的元素取决于配置（[EsmFold2Config]）和输入。

运行折叠主干：对输入进行特征化，将其嵌入到成对表示中，经过 num_loops 次回收迭代进行优化，然后读取距离图。这是结构预测的确定性部分；将返回的成对表示转换为 3D 坐标的扩散采样器位于 EsmFold2FoldingMixin 中——调用 [~EsmFold2Model.fold] 或 [~EsmFold2Model.infer_protein] 即可进行端到端预测。

distogram_logits (torch.Tensor，形状为 (batch_size, num_tokens, num_tokens, distogram_bins)) — 残基对上的预测距离分布 logits。
pair_states（形状为 (batch_size, num_tokens, num_tokens, pairwise_hidden_size) 的 torch.Tensor）— 主干网络最终的成对表示，以 fp32 格式存储。
single_inputs（形状为 (batch_size, num_tokens, single_inputs_size) 的 torch.Tensor）— 由输入嵌入器构建的拼接后的单输入特征。
relative_position_encoding（形状为 (batch_size, num_tokens, num_tokens, pairwise_hidden_size) 的 torch.Tensor）— 相对位置对编码。
token_bonds_encoding（形状为 (batch_size, num_tokens, num_tokens, pairwise_hidden_size) 的 torch.Tensor）— 嵌入的 token 间共价键特征。
atom_inputs (~models.esmfold2.modeling_esmfold2.EsmFold2AtomInputs，可选) — 经过特征化的参考构象原子输入，供扩散原子堆栈复用。

折叠

< source >

( attention_mask: Tensorasym_id: Tensormol_type: Tensordistogram_atom_idx: Tensornum_diffusion_samples: int | None = Nonenum_sampling_steps: int | None = None**trunk_kwargs )

端到端地从特征化输入预测结构：运行主干网络（EsmFold2Model.forward，其中记录了特征参数），从扩散结构头采样坐标，并使用置信度头对其进行评分。

此处仅列出了采样器和置信度头所需的参数；其余参数会原样转发至主干网络，因此 fold 接受的参数与 forward 相同。

attention_mask (torch.Tensor，形状为 (batch_size, num_tokens))：用于标记有效 token（1）与填充（0）的掩码。该参数也会转发给 trunk。 asym_id (torch.Tensor，形状为 (batch_size, num_tokens))：每个 token 的非对称单元（链）ID。该参数也会转发给 trunk。 mol_type (torch.Tensor，形状为 (batch_size, num_tokens))：每个 token 的分子类型编码：0 = 蛋白质，1 = DNA，2 = RNA，3 = 非聚合物（配体）。该参数也会转发给 trunk。 distogram_atom_idx (torch.Tensor，形状为 (batch_size, num_tokens))：每个 token 的代表原子（Cβ，甘氨酸为 Cα）的索引。供置信度头使用；trunk 不需要此参数。 num_diffusion_samples (int，可选)：要采样的并行结构样本数量；置信度头会为每个样本重新运行一次。默认值为 config.structure_head.num_diffusion_samples。 num_sampling_steps (int，可选)：扩散采样步数。默认值为 config.structure_head.inference_num_steps。 trunk_kwargs：其余的特征化输入（包括 num_loops，以及下方打包到 atom_inputs 中的原始逐原子张量），会转发给 [EsmFold2Model.forward]，该函数对这些参数有详细说明。

infer_protein

< source >

( sequence: str**forward_kwargs )

对 sequence 进行特征化并折叠。forward_kwargs 会传递到 [~EsmFold2Model.fold]。

infer_protein_as_pdb

< source >

( sequence: strsample_idx: int | None = None**forward_kwargs )

折叠 sequence 并将预测结果渲染为 PDB 字符串。

sample_idx 用于选择要渲染的扩散样本；默认情况下，选择排名最高的样本。

EsmFold2Output

类 transformers.models.esmfold2.modeling_esmfold2.EsmFold2Output

< source >

( distogram_logits: typing.Optional[torch.Tensor] = Nonesample_atom_coords: typing.Optional[torch.Tensor] = Noneplddt_logits: typing.Optional[torch.Tensor] = Noneplddt: typing.Optional[torch.Tensor] = Noneplddt_per_atom: typing.Optional[torch.Tensor] = Noneplddt_ca: typing.Optional[torch.Tensor] = Nonecomplex_plddt: typing.Optional[torch.Tensor] = Nonecomplex_iplddt: typing.Optional[torch.Tensor] = Nonepae_logits: typing.Optional[torch.Tensor] = Nonepae: typing.Optional[torch.Tensor] = Nonepde_logits: typing.Optional[torch.Tensor] = Nonepde: typing.Optional[torch.Tensor] = Noneresolved_logits: typing.Optional[torch.Tensor] = Noneptm: typing.Optional[torch.Tensor] = Noneiptm: typing.Optional[torch.Tensor] = Nonepair_chains_iptm: typing.Optional[torch.Tensor] = None )

参数

distogram_logits (torch.Tensor，形状为 (batch_size, num_tokens, num_tokens, distogram_bins)) — 残基对上的预测距离分布 logits（与 RNG 无关；不进行扩散采样）。
sample_atom_coords (torch.Tensor，形状为 (num_diffusion_samples, num_atoms, 3)) — 每个扩散样本预测的全原子笛卡尔坐标。
plddt_logits (torch.Tensor，形状为 (num_diffusion_samples, num_atoms, num_plddt_bins)) — 每个原子的 pLDDT 分箱 logits。
plddt (torch.Tensor，形状为 (num_diffusion_samples, num_tokens)) — 每个残基的预测 lDDT 置信度。
plddt_per_atom (torch.Tensor，形状为 (num_diffusion_samples, num_atoms)) — 每个原子的预测 lDDT 置信度。
plddt_ca (torch.Tensor，形状为 (num_diffusion_samples, num_tokens)) — 每个 token 的代表性（Cα）原子的预测 lDDT。
complex_plddt（形状为 (num_diffusion_samples,) 的 torch.Tensor）— 复合物所有原子的平均 pLDDT。
complex_iplddt (torch.Tensor，形状为 (num_diffusion_samples,)) — 界面加权复数 pLDDT。
pae_logits (torch.Tensor，形状为 (num_diffusion_samples, num_tokens, num_tokens, num_pae_bins)) — 预测对齐误差的 bin logits。
pae (torch.Tensor，形状为 (num_diffusion_samples, num_tokens, num_tokens)) — 每个残基对的预期预测对齐误差（Å）。
pde_logits (torch.Tensor，形状为 (num_diffusion_samples, num_tokens, num_tokens, num_pde_bins)) — 预测距离误差分箱的 logits。
pde (torch.Tensor，形状为 (num_diffusion_samples, num_tokens, num_tokens)) — 每个残基对的预期预测距离误差（Å）。
resolved_logits (torch.Tensor，形状为 (num_diffusion_samples, num_atoms, 2)) — 每个原子的已解析/未解析 logits。
ptm（形状为 (num_diffusion_samples,) 的 torch.Tensor）— 每个样本的预测 TM-score。
iptm（形状为 (num_diffusion_samples,) 的 torch.Tensor）— 每个样本的预测界面 TM 分数。
pair_chains_iptm (torch.Tensor，形状为 (num_diffusion_samples, num_chains, num_chains)) — 每个有序链对的预测界面 TM-score。

EsmFold2Model 的输出。所有置信度分数均在 0-1 范围内；每个样本的张量都有一个前置的 num_diffusion_samples 维度。

EsmFold2TrunkOutput

class transformers.models.esmfold2.modeling_esmfold2.EsmFold2TrunkOutput

< source >

( distogram_logits: typing.Optional[torch.Tensor] = Nonepair_states: typing.Optional[torch.Tensor] = Nonesingle_inputs: typing.Optional[torch.Tensor] = Nonerelative_position_encoding: typing.Optional[torch.Tensor] = Nonetoken_bonds_encoding: typing.Optional[torch.Tensor] = Noneatom_inputs: transformers.models.esmfold2.modeling_esmfold2.EsmFold2AtomInputs | None = None )

参数

distogram_logits (torch.Tensor，形状为 (batch_size, num_tokens, num_tokens, distogram_bins)) — 残基对上的预测距离分布 logits。
pair_states (torch.Tensor，形状为 (batch_size, num_tokens, num_tokens, pairwise_hidden_size)) — 主干的最终成对表示，以 fp32 格式存储。
single_inputs（形状为 (batch_size, num_tokens, single_inputs_size) 的 torch.Tensor）— 由输入嵌入器构建的拼接后的单输入特征。
relative_position_encoding (torch.Tensor，形状为 (batch_size, num_tokens, num_tokens, pairwise_hidden_size)) — 相对位置对编码。
token_bonds_encoding (torch.Tensor，形状为 (batch_size, num_tokens, num_tokens, pairwise_hidden_size)) — 嵌入的 token 间共价键特征。
atom_inputs (~models.esmfold2.modeling_esmfold2.EsmFold2AtomInputs，可选) — 经过特征化的参考构象原子输入，供扩散原子堆栈复用。

EsmFold2Model.forward() 的输出：折叠主干的成对表示、从中读取的 distogram，以及结构头和置信度头所使用的条件张量。除主干的随机初始成对状态外，给定输入后此处所有内容均为确定性的。

EsmFold2AtomInputs

class transformers.models.esmfold2.modeling_esmfold2.EsmFold2AtomInputs

< source >

( ref_pos: Tensorref_charge: Tensoratom_attention_mask: Tensorref_element: Tensorref_atom_name_chars: Tensorref_space_uid: Tensoratom_to_token: Tensor )

参数

ref_pos — (batch_size, num_atoms, 3) 参考构象的笛卡尔坐标。
ref_charge — (batch_size, num_atoms) 每个原子的形式电荷。
atom_attention_mask — (batch_size, num_atoms) 有效原子（1）与填充（0）。
ref_element — (batch_size, num_atoms, max_atomic_number) 原子序数的 one-hot 编码。
ref_atom_name_chars — (batch_size, num_atoms, max_chars, char_vocab_size) 原子名称的 one-hot 编码。
ref_space_uid — (batch_size, num_atoms) 每个原子的组 ID（即该原子的 token 索引）。
atom_to_token — (batch_size, num_atoms) 每个原子所属 token 的索引。

经过特征化的参考构象原子输入，打包在一起，以便原子堆栈只需一个参数。

作为 [EsmFold2Model.forward] 的一个参数传入，包含原始的 featurizer 张量。随后，EsmFold2Model._prepare_features 会对分类字段进行 one-hot 编码并将填充位置零，forward 则通过 dataclasses.replace 将这些编码后的值替换进去，因此该点之后的所有下游使用者看到的都是下文所述的编码形式。该结果在 [EsmFold2TrunkOutput] 中返回，以便扩散采样器和置信度头复用主干的 featurization，而无需重新计算。

使用普通 dataclass 而非 [~utils.ModelOutput]：此处的每个字段都是必需的，并且下游会无条件地读取它们，而 ModelOutput 最多只允许一个必需字段（其余字段必须默认为 None），这会将构造函数中缺失张量的错误转变为在原子编码器深处浮现的 None。

Update on GitHub

←ESMC EuroBERT→