
前两章我们介绍了ESM3和ESM-C模型
ESM3蛋白质语言模型cookbook(1)
ESM3蛋白质语言模型cookbook(2)

·Hayes论文图
0.背景
今天我们将会进行下一部分的介绍,首先是Hayes的论文:(
Simulating 500 million years of evolution with a language model
,10.1126/science.ads0018)
论文地址:https://www.science.org/doi/10.1126/science.ads0018
在这篇论文中, Hayes团队使用esm3模型进行了green fluorescent protein,GFP的生成。green fluorescent protein是绿色荧光蛋白质的全称,荧光蛋白是水母和珊瑚发光的主要原因,并且也是现代生物技术的重要工具。它们有一个“优雅”的结构:一个十一边形的β桶,其中心有螺旋结构,它为蛋白质自身原子形成发光发色团提供了支架。
这种机制在本质上是独一无二的——没有其他蛋白质能自发地从自己的结构中形成荧光发色团——这表明,即使对自然界来说,产生荧光也很困难。
本次教程介绍了 Hayes 等人用于设计 esmGFP 的计算方法。尽管 esmGFP 与已知荧光蛋白的序列相似性较低(仅有 58% 的一致性),但它的亮度和光谱特性与自然界中的 GFP 非常相似。此外,我们还发现了许多其他亮度高、序列一致性相似或更高的新型 GFP。利用本笔记本中概述的方法,人们很可能可以设计出更多新的 GFP。
这本笔记本实现了用于启动生成 esmGFP 构思过程的核心提示词(prompt)。其中实际使用的整体流程有两个主要不同之处:
废话不多说,直接开始今天的教程吧:
1.环境设置
from IPython.display import clear_output
!pip install git+https://github.com/evolutionaryscale/esm.git
!pip install py3Dmol
clear_output() # Suppress pip install log lines after installation is complete.from getpass import getpass
import biotite.sequence as seq
import biotite.sequence.align as align
import biotite.sequence.graphics as graphics
import matplotlib.pyplot as pl
import py3Dmol
import torch
from esm.sdk import client
from esm.sdk.api import ESMProtein, GenerationConfig
from esm.utils.structure.protein_chain import ProteinChain继续使用ems3的api进行操作:
token = getpass("Token from Forge console: ")
然后我们创建一个模型,它的行为有点像一个 PyTorch 模型,但实际上它会将输入发送到 Forge 服务器,在远程服务器上通过神经网络权重运行这些输入,然后将输出张量返回到这个笔记本中。这个桩对象也可以在 EvolutionaryScale 的 SDK 中使用,以简化生成、折叠以及一般采样相关的许多操作。这一点很重要,因为迭代采样是获得 ESM3 最佳性能的关键,而 SDK 管理了实现这些标准流程所涉及的大量复杂性。
model = client(
model="esm3-medium-2024-03", url="https://forge.evolutionaryscale.ai", token=token
)2.构造GFP提示词
ESM3 是一个生成模型。要使用其生成能力,我们需要熟悉如何构建提示词(prompts)。ESM3 能够在蛋白质的序列、结构和功能之间进行联合推理,因此我们可以构建新的提示方式,以比许多其他生物语言模型更高的控制水平引导模型生成蛋白质。
序列、结构和功能这三种模态在模型中被表示为离散 token 的轨道,既出现在模型的输入中,也出现在输出中,并在模型内部融合为一个统一的潜在空间。ESM3 使用带有可变掩码率的生成式掩码语言建模目标进行训练,因此我们可以通过完全或部分掩码的上下文,以及在不同轨道上的不同条件点进行提示。这使我们在提示词的设定上具有极大的创造空间。
提示工程既是一门艺术,也是一门科学,因此通常需要通过实验来获得能生成期望结果的提示词。此外,由于我们是通过采样方式从模型中生成内容,同一个提示词的不同次生成结果会有所不同。有些提示词成功率较高,只需生成几次就能获得一个候选蛋白设计;而另一些更复杂的提示词可能需要生成上千次!通过对模型进行对齐训练,可以使其更具可控性。
我们将使用的是未经对齐的原始预训练模型,但我们已经对这个提示做了大量优化,因此通常只需少量生成就能获得一个有趣的设计。
我们将从 PDB 数据库中的 1qy3 序列和结构片段构建提示。以下代码从 PDB 获取数据,然后使用 ESM3 的分词器将序列和结构转换为可以传入模型的 token。可以看到,氨基酸类型和空间坐标都被转换为每个序列位置一个离散 token。
template_gfp = ESMProtein.from_protein_chain(
ProteinChain.from_rcsb("1qy3", chain_id="A")
)
template_gfp_tokens = model.encode(template_gfp)
print("Sequence tokens:")
print(
" ", ", ".join([str(token) for token in template_gfp_tokens.sequence.tolist()])
)
print("Structure tokens:")
print(
" ", ", ".join([str(token) for token in template_gfp_tokens.structure.tolist()])
)这段代码的作用是:从pdb数据中加载一个GFP蛋白质结构(pdb id:1qy3,链:A),并且用esm3模型的编码器将其转化为序列和结构的token表示,用于后续提示词的构建。
输出:

接下来:我们现在将构建一个提示。具体来说,我们将指定在靠近我们希望形成发色团的位置上的4个氨基酸种类,以及在β桶结构上已知有助于发色团形成的2个氨基酸种类。
此外,我们将通过将编码的1qy3结构中的标记添加到提示中的结构轨道,在这些位置上指定其结构应与1qy3结构相似。我们还将另外指定几个位置(沿着α螺旋的弯折处)。
prompt_sequence = ["_"] * len(template_gfp.sequence)
prompt_sequence[59] = "T"
prompt_sequence[62] = "T"
prompt_sequence[63] = "Y"
prompt_sequence[64] = "G"
prompt_sequence[93] = "R"
prompt_sequence[219] = "E"
prompt_sequence = "".join(prompt_sequence)
print(template_gfp.sequence)
print(prompt_sequence)
prompt = model.encode(ESMProtein(sequence=prompt_sequence))
prompt.structure = torch.full_like(prompt.sequence, 4096)
prompt.structure[0] = 4098
prompt.structure[-1] = 4097
prompt.structure[55:70] = template_gfp_tokens.structure[56:71]
prompt.structure[93] = template_gfp_tokens.structure[93]
prompt.structure[219] = template_gfp_tokens.structure[219]
print("".join(["✔" if st < 4096 else "_" for st in prompt.structure]))输出:

输出显示了原始的1qy3序列,以及我们在提示中设定的序列轨道上的氨基酸种类,以及在结构轨道上设置了标记的位置。ESM3接下来的任务是补全剩余被掩盖(下划线)的结构和序列。
有一个小说明:我们在提示中引入了突变 A93:R。这并不是错误。将该位置设为丙氨酸(Alanine)会导致发色团成熟得非常缓慢(这正是我们能够测量GFP环化前结构的原因)。不过,我们并不想等待GFP慢慢发光,因此我们选择了精氨酸(Arginine)来替代这个位置。
3.结构生成
接着我们向模型提供提示词,并对结构标记轨道进行解码。这类似于为一个活性位点提示创建主链骨架,但也存在一些细微差别。例如,由于我们已经在提示中指定了部分结构标记(例如活性位点及其关键对应残基附近),模型实际上会围绕这些已有结构生成其余部分。
ESM3会以迭代方式采样生成标记。标记可以逐个采样,也可以并行采样,顺序不限,直到所有位置都被完整生成。EvolutionaryScale SDK 中的 generate() 函数实现了一种我们认为在从模型中采样时较为有效的方法。
%%time
num_tokens_to_decode = min((prompt.structure == 4096).sum().item(), 20)
structure_generation = model.generate(
prompt,
GenerationConfig(
# Generate a structure.
track="structure",
# Sample one token per forward pass of the model.
num_steps=num_tokens_to_decode,
# Sampling temperature trades perplexity with diversity.
temperature=1.0,
),
)
print("These are the structure tokens corresponding to our new design:")
print(
" ", ", ".join([str(token) for token in structure_generation.structure.tolist()])
)
# Decodes structure tokens to backbone coordinates.
structure_generation_protein = model.decode(structure_generation)输出:

现在让我们可视化生成的结构。这很可能会呈现出一个熟悉的GFP结构:围绕α螺旋形成的β桶。
view = py3Dmol.view(width=1000, height=500)
view.addModel(
structure_generation_protein.to_protein_chain().infer_oxygen().to_pdb_string(),
"pdb",
)
view.setStyle({"cartoon": {"color": "lightgreen"}})
view.zoomTo()
view.show()
·通过esm3生成的类gfp蛋白质结构
接下来我们再去看看pdb id为1qy3 的蛋白质结构是什么样子的呢?

·pdb官网的1qy3的结构
此时,只有在以下条件满足时,我们才会继续生成:
设计在活性位点上与野生型GFP高度匹配;整个蛋白结构上有一定差异(否则就会和野生型GFP的序列非常接近);整体结构仍像经典的GFP,即α螺旋被β桶包围。
当然,如果要生成很多设计,我们无法逐个手动查看,因此我们采用自动的拒绝采样标准,主要根据整体结构的RMSD和受限位点的RMSD,判断生成的结构是否忠实于最初的提示信息。如果这些检查通过,就可以尝试为这个结构设计相应的序列;如果不通过,就回到前面几步,重新生成结构直到通过这些计算筛选。当然也可以不这样做,毕竟这就是你自己设计的GFP。
constrained_site_positions = [59, 62, 63, 64, 93, 219]
template_chain = template_gfp.to_protein_chain()
generation_chain = structure_generation_protein.to_protein_chain()
constrained_site_rmsd = template_chain[constrained_site_positions].rmsd(
generation_chain[constrained_site_positions]
)
backbone_rmsd = template_chain.rmsd(generation_chain)
c_pass = "✅" if constrained_site_rmsd < 1.5 else "❌"
b_pass = "✅" if backbone_rmsd > 1.5 else "❌"
print(f"Constrained site RMSD: {constrained_site_rmsd:.2f} Ang {c_pass}")
print(f"Backbone RMSD: {backbone_rmsd:.2f} Ang {b_pass}")输出:

上面的代码比较了新生成的蛋白结构与模板gfp在特定位置和整体结构上的差异。
4.序列设计
现在我们已经有了一个具有一定结构变化的骨架,同时它在GFP的受限位点上也保持一致。接下来我们希望设计一个能折叠成该结构的氨基酸序列。我们可以使用上一次的生成结果(即原始提示加上新的结构标记)作为新的提示,再次调用ESM3。
一旦设计出序列,我们还需要确认这个序列是否确实能折叠成我们期望的结构。为此,我们会从提示中移除所有其他条件,仅使用该序列进行折叠。在ESM3中,这个过程非常方便,只需根据氨基酸序列生成一组结构标记即可。此时我们希望模型给出最有信心的预测结果(不引入多样性),因此会将采样温度设为零(temperature=0.0)。
%%time
# Based on internal research, there's not a benefit to iterative decoding past 20 steps
num_tokens_to_decode = min((prompt.sequence == 32).sum().item(), 20)
sequence_generation = model.generate(
# Generate a sequence.
structure_generation,
GenerationConfig(track="sequence", num_steps=num_tokens_to_decode, temperature=1.0),
)
# Refold
sequence_generation.structure = None
length_of_sequence = sequence_generation.sequence.numel() - 2
sequence_generation = model.generate(
sequence_generation,
GenerationConfig(track="structure", num_steps=1, temperature=0.0),
)
# Decode to AA string and coordinates.
sequence_generation_protein = model.decode(sequence_generation)输出:

从esmprotein类中得到序列:
sequence_generation_protein.sequence输出:

我们可以将这个序列与原始模板进行比对,看看它与avGFP有多相似。同时,也可以将其与所有已知的荧光蛋白进行比对,以评估这个潜在GFP的新颖性。
seq1 = seq.ProteinSequence(template_gfp.sequence)
seq2 = seq.ProteinSequence(sequence_generation_protein.sequence)
alignments = align.align_optimal(
seq1, seq2, align.SubstitutionMatrix.std_protein_matrix(), gap_penalty=(-10, -1)
)
alignment = alignments[0]
identity = align.get_sequence_identity(alignment)
print(f"Sequence identity: {100*identity:.2f}%")
print("\nSequence alignment:")
fig = pl.figure(figsize=(8.0, 4.0))
ax = fig.add_subplot(111)
graphics.plot_alignment_similarity_based(
ax, alignment, symbols_per_line=45, spacing=2, show_numbers=True
)
fig.tight_layout()
pl.show()输出:

可以看到序列的同一性是33.76%
我们现在重新检查受限位点的计算指标。如果发现这些位点不匹配,我们就需要重新尝试设计这个序列。如果多次尝试都无法设计出一个与结构相匹配的序列,那很可能说明这个结构本身不容易设计出稳定的序列,这时候我们也可以考虑放弃这个结构。
此时,骨架的RMSD就不是特别重要了,只要最终生成的序列与原始GFP有足够的差异,能够满足我们对科学探索的兴趣就可以。
template_chain = template_gfp.to_protein_chain()
generation_chain = sequence_generation_protein.to_protein_chain()
constrained_site_rmsd = template_chain[constrained_site_positions].rmsd(
generation_chain[constrained_site_positions]
)
backbone_rmsd = template_chain.rmsd(generation_chain)
c_pass = "✅" if constrained_site_rmsd < 1.5 else "❌"
b_pass = "🤷♂️"
print(f"Constrained site RMSD: {constrained_site_rmsd:.2f} Ang {c_pass}")
print(f"Backbone RMSD: {backbone_rmsd:.2f} Ang {b_pass}")输出:

接下来再做张图来看看:
view = py3Dmol.view(width=600, height=600)
view.addModel(sequence_generation_protein.to_pdb_string(), "pdb")
view.setStyle({"cartoon": {"color": "lightgreen"}})
view.zoomTo()
view.show()输出:

·通过序列得到的结构
5.总结
Finally,我们本章cookbook实现了从token层面的蛋白质序列设计,其中背后的引擎则是esm3蛋白质多模态模型,并且也用到了自然语言中的很多概念,如mask和token的表示。大家可以自己运行一下代码,有问题欢迎一起交流~
您的点赞是我更新的动力,谢谢~