在 Amazon SageMaker 上使用 ESMFold 语言模型加速蛋白质结构预测

蛋白质驱动着许多生物过程，如酶活性、分子输运和细胞支持。通过蛋白质的三维结构，可以深入了解蛋白质的功能以及蛋白质如何与其他生物分子相互作用。测定蛋白质结构的实验方法（如 X 射线晶体学和核磁共振波谱学）既昂贵又耗时。相比之下，最近开发的几种计算方法能够根据蛋白质的氨基酸序列快速准确地预测蛋白质的结构。这些方法对于难以通过实验研究的蛋白质至关重要，例如膜蛋白，它是许多药物的靶标。

在这些方法中，一个众所周知的例子是 AlphaFold，这是一种基于深度学习的算法，因其预测准确而闻名于世；来自 FAIR 的 ESMFold 是另一种基于深度学习的高精度方法，用于根据氨基酸序列预测蛋白质结构。ESMFold 以大型蛋白质语言模型（pLM）为骨干，进行端到端的操作。与 AlphaFold2 不同，这种方法不需要查找或多序列比对（MSA）步骤，也不依赖外部数据库来生成预测结果。取而代之的是，FAIR Research 团队根据 UniRef 的数百万个蛋白质序列对模型进行了训练。在训练过程中，该模型开发了注意力模式，这些模式优雅地代表了序列中氨基酸之间的进化相互作用。使用 pLM 代替 MSA 后，可使预测时间比其他先进模型快60倍。这种速度使 FAIR 团队能够创建开源的 ESM 宏基因组图谱，该图谱由超过7.7亿个宏基因组蛋白质结构预测组成。

在本文中，我们使用 Hugging Face Hub 上提供的预训练 ESMFold 模型和 Amazon SageMaker，来预测曲妥珠单抗的重链结构，曲妥珠单抗是一种由 Genentech 首次开发的用于治疗 HER2 阳性乳腺癌的单克隆抗体。如果研究人员想测试序列修饰的效果，快速预测这种蛋白质的结构可能非常有用。这有可能提高患者的生存率或减少药物副作用。这篇文章提供了一个 Jupyter notebook 示例，相关脚本位于以下 GitHub 存储库中：https://github.com/aws-samples/aws-healthcare-lifescience-ai-ml-sample-notebooks/tree/main/workshops/AI_Driven_Protein_Analysis

先决条件

我们建议在 Amazon SageMaker 工作室笔记本电脑中运行这个示例，该笔记本电脑在 ml.r5.xlarge 实例类型上运行 PyTorch 1.13 Python 3.9 CPU 优化的映像。

以视觉形式呈现

曲妥珠单抗的实验结构

首先，我们使用 biopython 库和辅助脚本从 RCSB 蛋白质数据库下载曲妥珠单抗结构：

from Bio.PDB import PDBList, MMCIFParser
from prothelpers.structure import atoms_to_pdb


target_id = "1N8Z"
pdbl = PDBList()
filename = pdbl.retrieve_pdb_file(target_id, pdir="data")
parser = MMCIFParser()
structure = parser.get_structure(target_id, filename)
pdb_string = atoms_to_pdb(structure)

左滑查看更多

接下来，我们使用 py3Dmol 库将结构可视化为交互式 3D 视觉内容：

view = py3Dmol.view()
view.addModel(pdb_string)
view.setStyle({'chain':'A'},{"cartoon": {'color': 'orange'}})
view.setStyle({'chain':'B'},{"cartoon": {'color': 'blue'}})
view.setStyle({'chain':'C'},{"cartoon": {'color': 'green'}})
view.show()

左滑查看更多

下图显示了蛋白质数据库（PDB）中 1N8Z 的 3D 蛋白质结构。在此图中，曲妥珠单抗轻链显示为橙色，重链显示为蓝色（可变区为浅蓝色），HER2 抗原显示为绿色。

我们将首先使用 ESMFold，根据氨基酸序列预测重链（链 B）的结构。然后，将预测结果与上面显示的实验测定的结构进行对比。

使用 ESMFold

根据曲妥珠单抗重链序列预测其结构

让我们使用 ESMFold 模型来预测重链的结构，并将其与实验结果进行对比。首先，使用 Studio 中预置的 notebook 环境，该环境预装了几个重要的库，如 PyTorch。虽然我们能够使用加速实例类型来提高 notebook 分析的性能，但现在我们改用非加速实例，并在 CPU 上运行 ESMFold 预测。

首先，从 Hugging Face Hub 加载预训练的 ESMFold 模型和标记器：

from transformers import AutoTokenizer, EsmForProteinFolding


tokenizer = AutoTokenizer.from_pretrained("facebook/esmfold_v1")
model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_usage=True)

左滑查看更多

接下来，将模型复制到我们的设备（本例中为 CPU），并设置一些模型参数：

device = torch.device("cpu")
model.esm = model.esm.float()
model = model.to(device)
model.trunk.set_chunk_size(64)

左滑查看更多

为了准备用于分析的蛋白质序列，我们需要对其进行标记化处理。这样可以将氨基酸符号（EVQLV…）转换为 ESMFold 模型可以理解的数字格式（6,19,5,10,19,…）：

tokenized_input = tokenizer([experimental_sequence], return_tensors="pt", add_special_tokens=False)["input_ids"]
tokenized_input = tokenized_input.to(device)

左滑查看更多

接下来，我们将标记化输入复制到模式中，进行预测，并将结果保存到文件中：

with torch.no_grad():
notebook_prediction = model.infer_pdb(experimental_sequence)
with open("data/prediction.pdb", "w") as f:
f.write(notebook_prediction)

左滑查看更多

对于非加速实例类型（如 r5），这大约需要3分钟。

我们可以通过对比实验结构，来检验 ESMFold 预测的准确性。我们使用密歇根大学 Zhang Lab 开发的 US-Align 工具来完成这项工作：

from prothelpers.usalign import tmscore


tmscore("data/prediction.pdb", "data/experimental.pdb", pymol="data/superimposed")

左滑查看更多

模板建模得分（TM-score）是评测蛋白质结构相似性的指标。得分为 1.0 表示完全匹配。得分高于 0.7 表示两个蛋白质具有相同的骨干结构。得分高于 0.9 表示蛋白质在下游使用中具有功能互换性。在我们的示例中，TM-Score 达到 0.802，表示 ESMFold 预测可能适用于结构评分或配体结合实验等应用，但可能不适用于分子置换等需要极高精度的使用案例。

我们可以通过可视化对齐结构来验证这一结果。这两个结构的重叠程度很高，但并不完全重叠。蛋白质结构预测是一个快速发展的领域，许多研究团队都在开发越来越精确的算法！

将 ESMFold 部署为

SageMaker 推理端点

在 notebook 中运行模型推理可用于实验目的，但如果需要将模型与应用程序集成呢？或者是与 MLOps 管道集成呢？在这种情况下，更好的选择是将模型部署为推理端点。在下面的示例中，我们将 ESMFold 作为 SageMaker 实时推理端点部署在加速实例上。SageMaker 实时端点提供了一种可扩展、经济高效且安全的方式来部署和托管机器学习（ML）模型。通过自动扩缩，您可以调整运行端点的实例数量以满足应用程序的需求，从而优化成本并确保高可用性。

利用预置的 Hugging Face 的 SageMaker 容器，可以轻松地为常见任务部署深度学习模型。然而，对于像蛋白质结构预测这样的新使用案例，我们需要定义一个自定义的 inference.py 脚本，来加载模型、运行预测和格式化输出内容。该脚本包含的代码与我们在 notebook 中使用的代码基本相同。我们还创建了一个 requirements.txt 文件来定义一些 Python 依赖关系，供我们的端点使用。您可以在 GitHub 存储库中看到我们创建的文件：https://github.com/aws-samples/aws-healthcare-lifescience-ai-ml-sample-notebooks/tree/main/workshops/AI_Driven_Protein_Analysis

在下图中，曲妥珠单抗重链的实验（蓝色）和预测（红色）结构非常相似，但并不完全相同。

在 code 目录中创建了必要的文件后，我们使用 SageMaker HuggingFaceModel 类部署模型。该类使用一个预置的容器，简化将 Hugging Face 模型部署到 SageMaker 的过程。请注意，创建端点可能需要 10 分钟或更长时间，具体取决于我们区域中 ml.g4dn 实例类型的可用性。

from sagemaker.huggingface import HuggingFaceModel
from datetime import datetime


huggingface_model = HuggingFaceModel(
model_data = model_artifact_s3_uri, # Previously staged in S3
name = f"emsfold-v1-model-" + datetime.now().strftime("%Y%m%d%s"),
transformers_version='4.17',
pytorch_version='1.10',
py_version='py38',
role=role,
source_dir = "code",
entry_point = "inference.py"
)


rt_predictor = huggingface_model.deploy(
initial_instance_count = 1,
instance_type="ml.g4dn.2xlarge",
endpoint_name=f"my-esmfold-endpoint",
serializer = sagemaker.serializers.JSONSerializer(),
deserializer = sagemaker.deserializers.JSONDeserializer()
)

左滑查看更多

完成端点部署后，我们可以重新提交蛋白质序列，并显示预测的前几行内容：

endpoint_prediction = rt_predictor.predict(experimental_sequence)[0]
print(endpoint_prediction[:900])

左滑查看更多

由于我们将端点部署到加速实例中，因此预测应该只需几秒钟。结果中的每一行对应于一个原子，包括氨基酸标识、三个空间坐标和表示该位置预测置信度的 pLDDT 分数。

使用与以前相同的方法，我们可以看到 notebook 和端点的预测结果完全相同。

如下图所示，在 notebook 中生成的 ESMFold 预测结果（红色）与端点生成的预测结果（蓝色）完全一致。

清理

为了避免产生进一步的费用，我们删除了推理端点和测试数据：

rt_predictor.delete_endpoint()
bucket = boto_session.resource("s3").Bucket(bucket)
bucket.objects.filter(Prefix=prefix).delete()
os.system("rm -rf data obsolete code")

左滑查看更多

小结

蛋白质结构计算预测是了解蛋白质功能的重要工具。除基础研究外，AlphaFold 和 ESMFold 等算法在医学和生物技术领域也有很多应用。利用这些模型生成的结构洞察，有助于我们更好地了解生物分子是如何相互作用的。这样就能为患者提供更好的诊断工具和疗法。

在这篇文章中，我们将展示如何使用 SageMaker 将 Hugging Face Hub 的 ESMFold 蛋白质语言模型部署为可扩展的推理端点。有关在 SageMaker 上部署 Hugging Face 模型的更多信息，请参阅结合使用 Hugging Face 与 Amazon SageMaker (https://docs.aws.amazon.com/sagemaker/latest/dg/hugging-face.html)。您还可以在 Awesome Protein Analysis on Amazon GitHub 存储库中找到更多蛋白质科学示例。

Original URL:

https://aws.amazon.com/blogs/machine-learning/accelerate-protein-structure-prediction-with-the-esmfold-language-model-on-amazon-sagemaker/

本篇作者