ESM3蛋白质语言模型cookbook（1）

Tom2Code

发布于 2026-04-17 17:25:04

290

这次即将推出一个ESM3的系列教程来展示ESM3的用法。

一.背景介绍

对于蛋白质序列的研究和结构的研究现在变得越来越热门，所谓工欲善其事，必先利其器，所以今天我们就来介绍一下ESM3, 使用ESM3可以对蛋白质进行特征提取，无论是之后再接入各种各样的神经网络，都是很快速的一种方法。例如，下面介绍到的蛋白质坐标的提取可以直接用于结构的预测或者结构的生成（RFDiffusion的范式）。

图中为Esm的官网 www.evolutionaryscale.ai 读者需要自行注册

ESM的工作台

ESM-2 是 Meta AI 推出的蛋白质语言模型，基于约 6500 万条天然序列训练的大型 Transformer，可直接从一级序列预测近原子分辨率的三维结构，性能与 AlphaFold 旗鼓相当。

hugggingface上ems2的model card

紧接着，原团队核心成员成立 EvolutionaryScale 并发布生成式旗舰 ESM-3——参数规模再跃升、支持自然语言式“提示”，能模拟 5 亿年进化并实测设计出全新的绿色荧光蛋白，帮助科学家按需创造酶、抗体及碳捕获蛋白等功能分子，标志着蛋白质语言模型从“预测”迈向“可控生成”时代。ESM-3更是一个多模态的蛋白质语言模型，结合了蛋白质的结构和功能等特征。

二.ESMProtein类介绍

2.1前言

在本节内容中，我们将熟悉ESMProtein类，该类包含代表序列、结构和功能的蛋白质的多种属性。ESM3模型使用输入（提示）中的这些属性，并将其作为输出的一部分生成。

接下来解释一下这五个属性：

sequence序列：氨基酸序列

coordinates坐标：蛋白质每个氨基酸中原子的三维坐标

secondary_structure二级结构：8级二级结构（SS8）

sasa：溶剂可及表面积（sasa）

function_annotations：从InterPro(https://www.ebi.ac.uk/interpro/)派生的函数注释

2.2依赖导入

废话不多讲，还是通过代码来讲解案例，接下来就是安装和导入依赖的代码说明：

# Install esm and other dependencies
! pip install esm
! pip install py3Dmol
! pip install matplotlib
! pip install dna-features-viewer

2.3创建ESMProtein对象

from biotite.database import rcsb
from esm.sdk.api import ESMProtein
from esm.utils.structure.protein_chain import ProteinChain
from esm.utils.types import FunctionAnnotation
pdb_id = "1cm4"
chain_id = "A"

str_io = rcsb.fetch(pdb_id, "pdb")
protein_chain = ProteinChain.from_pdb(str_io, chain_id=chain_id, id=pdb_id)
protein = ESMProtein.from_protein_chain(protein_chain)

上述代码示例展示了一个规范的蛋白质结构加载流程：首先借助 Biotite 的 rcsb.fetch 接口，从 RCSB Protein Data Bank 在线检索并下载指定 PDB 条目（此处为 1CM4 的 A 链）；随后使用 ProteinChain.from_pdb 将获取到的 PDB 字符串解析为 ProteinChain 对象，并显式保留链编号与 PDB ID，确保后续索引与注释的一致性；最后通过 ESM SDK 提供的 ESMProtein.from_protein_chain 方法，将该链封装为符合 ESM 系列模型输入规范的 ESMProtein 实例，为结构预测、功能注释或定向设计等下游任务奠定标准化数据基础。

来打印一下1CM4这个蛋白质的sequence吧：

接下来我们再加载一下这个蛋白质的坐标看看吧：

print(protein.coordinates.shape) #先看看坐标的形状

yep，就是十分的方便，可以看到已经转化成为tensor的格式了，这样方便我们以后对其进行运算和处理。对这一组数字的解释：

143: 这里的143就是序列的长度

37: 37是氨基酸中原子的最大可能数量，以atom37表示。如果结构中不存在某些原子，它们将显示为nan。因为蛋白质由氨基酸组成，而每一种氨基酸又由不同的原子组成，所以这里是原子坐标，如果没有这个原子，那么其坐标则为nan

3: 每个原子的(x,y,z)坐标

打印一下看看咯：

print(protein.coordinates)

2.4 创建可视化函数

visualize_3D_coordinates函数是通过创建一个包含所有丙氨酸的pdb文件，直接从坐标张量进行可视化

visualize_3D_protein函数是从ESMProtein实例中可视化，该实例具有正确的氨基酸

import py3Dmol
def visualize_pdb(pdb_string):
    view = py3Dmol.view(width=400, height=400)
    view.addModel(pdb_string, "pdb")
    view.setStyle({"cartoon": {"color": "spectrum"}})
    view.zoomTo()
    view.render()
    view.center()
    return view
def visualize_3D_coordinates(coordinates):
    """
    This uses all Alanines
    """
    protein_with_same_coords = ESMProtein(coordinates=coordinates)
    # pdb with all alanines
    pdb_string = protein_with_same_coords.to_pdb_string()
    return visualize_pdb(pdb_string)
def visualize_3D_protein(protein):
    pdb_string = protein.to_pdb_string()
    return visualize_pdb(pdb_string)

这两个函数的核心差异：

接下来我们来看看这两个函数的可视化区别：

这个函数的输入是蛋白质序列的坐标

visualize_3D_coordinates(protein.coordinates)

接下来是第二个可视化函数的结果：

这个输入则是一个完整的ESMProtein对象

visualize_3D_protein(protein)

结构肯定是一样的，只不过是我们的输入的序列信息不同。

2.5二级结构的统计

secondary_structure属性包含二级结构的表示。在高级别的分类中，我们可以将每个氨基酸分为三类：α螺旋、β折叠和卷曲，我们可以在前面的3D可视化中看到。

ESMProtein使用8类二级结构，可以在给定三维原子坐标的情况下用dssp计算。由于安装dssp和安装esm包是一个单独的过程，在本笔记本中，我们将展示如何使用黑云母的annotate_sse计算更粗略的3类分类。我们可以用这个3类分类设置secondary_structure属性。

这里我们使用一个函数统计二级结构信息：

from biotite.structure import annotate_sse
def get_approximate_ss(protein_chain: ProteinChain):
    # get biotite's ss3 representation
    ss3_arr = annotate_sse(protein_chain.atom_array)
    biotite_ss3_str = "".join(ss3_arr)
    # translate into ESM3's representation
    translation_table = str.maketrans(
        {
            "a": "H",  # alpha helix
            "b": "E",  # beta sheet
            "c": "C",  # coil
        }
    )
    esm_ss3 = biotite_ss3_str.translate(translation_table)
    return esm_ss3

该函数接收 ProteinChain 对象，先用 Biotite 的 annotate_sse 推断每个残基的二级结构（α-螺旋 a、β-折叠 b、无规卷曲 c），再映射成 ESM-3 规范的 “H/E/C” 标记并按序拼接返回。

protein.secondary_structure = get_approximate_ss(protein_chain)
print(protein.secondary_structure)

输出：

这里的H,E,C分别代表了不同的二级结构。

接下来我们还是定义一个可视化的函数：

import biotite
import biotite.sequence as seq
import biotite.sequence.graphics as graphics
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Rectangle
# Create 'FeaturePlotter' subclasses
# for drawing the secondary structure features
class HelixPlotter(graphics.FeaturePlotter):
    def __init__(self):
        pass
    # Check whether this class is applicable for drawing a feature
    def matches(self, feature):
        if feature.key == "SecStr":
            if "sec_str_type" in feature.qual:
                if feature.qual["sec_str_type"] == "helix":
                    return True
        return False
    # The drawing function itself
    def draw(self, axes, feature, bbox, loc, style_param):
        # Approx. 1 turn per 3.6 residues to resemble natural helix
        n_turns = np.ceil((loc.last - loc.first + 1) / 3.6)
        x_val = np.linspace(0, n_turns * 2 * np.pi, 100)
        # Curve ranges from 0.3 to 0.7
        y_val = (-0.4 * np.sin(x_val) + 1) / 2
        # Transform values for correct location in feature map
        x_val *= bbox.width / (n_turns * 2 * np.pi)
        x_val += bbox.x0
        y_val *= bbox.height
        y_val += bbox.y0
        # Draw white background to overlay the guiding line
        background = Rectangle(
            bbox.p0, bbox.width, bbox.height, color="white", linewidth=0
        )
        axes.add_patch(background)
        axes.plot(x_val, y_val, linewidth=2, color=biotite.colors["dimgreen"])
class SheetPlotter(graphics.FeaturePlotter):
    def __init__(self, head_width=0.8, tail_width=0.5):
        self._head_width = head_width
        self._tail_width = tail_width
    def matches(self, feature):
        if feature.key == "SecStr":
            if "sec_str_type" in feature.qual:
                if feature.qual["sec_str_type"] == "sheet":
                    return True
        return False
    def draw(self, axes, feature, bbox, loc, style_param):
        x = bbox.x0
        y = bbox.y0 + bbox.height / 2
        dx = bbox.width
        dy = 0
        if loc.defect & seq.Location.Defect.MISS_RIGHT:
            # If the feature extends into the previous or next line
            # do not draw an arrow head
            draw_head = False
        else:
            draw_head = True
        axes.add_patch(
            biotite.AdaptiveFancyArrow(
                x,
                y,
                dx,
                dy,
                self._tail_width * bbox.height,
                self._head_width * bbox.height,
                # Create head with 90 degrees tip
                # -> head width/length ratio = 1/2
                head_ratio=0.5,
                draw_head=draw_head,
                color=biotite.colors["orange"],
                linewidth=0,
            )
        )
# Converter for the DSSP secondary structure elements
# to the classical ones
dssp_to_abc = {
    "I": "c",
    "S": "c",
    "H": "a",
    "E": "b",
    "G": "c",
    "B": "b",
    "T": "c",
    "C": "c",
}
def visualize_secondary_structure(sse, first_id):
    """
    Helper function to convert secondary structure array to annotation
    and visualize it.
    """
    def _add_sec_str(annotation, first, last, str_type):
        if str_type == "a":
            str_type = "helix"
        elif str_type == "b":
            str_type = "sheet"
        else:
            # coil
            return
        feature = seq.Feature(
            "SecStr", [seq.Location(first, last)], {"sec_str_type": str_type}
        )
        annotation.add_feature(feature)
    # Find the intervals for each secondary structure element
    # and add to annotation
    annotation = seq.Annotation()
    curr_sse = None
    curr_start = None
    for i in range(len(sse)):
        if curr_start is None:
            curr_start = i
            curr_sse = sse[i]
        else:
            if sse[i] != sse[i - 1]:
                _add_sec_str(
                    annotation, curr_start + first_id, i - 1 + first_id, curr_sse
                )
                curr_start = i
                curr_sse = sse[i]
    # Add last secondary structure element to annotation
    _add_sec_str(annotation, curr_start + first_id, i + first_id, curr_sse)
    fig = plt.figure(figsize=(30.0, 3.0))
    ax = fig.add_subplot(111)
    graphics.plot_feature_map(
        ax,
        annotation,
        symbols_per_line=150,
        loc_range=(first_id, first_id + len(sse)),
        feature_plotters=[HelixPlotter(), SheetPlotter()],
    )
    fig.tight_layout()
    return fig, ax
def plot_ss8(ss8_string):
    ss3 = np.array([dssp_to_abc[e] for e in ss8_string], dtype="U1")
    _, ax = visualize_secondary_structure(ss3, 1)
    ax.set_xticks([])

我们自定义了 HelixPlotter 与 SheetPlotter 两个 FeaturePlotter 子类，用专门的画笔在 Matplotlib 上分别以连续正弦曲线和带箭头的粗线渲染 α-螺旋和 β-折叠；随后将 DSSP 的 8-态符号映射为经典 “a/b/c” 三态，借助 Biotite 的 Annotation 机制把二级结构字符串定位到具体残基区间，最后通过 plot_feature_map() 在单幅图中直观展示整条蛋白序列的螺旋-折叠-线圈分布。

打印一下看看咯：

不同的曲线代表了不同的二级结构。

2.6 InterPro信息

接下来，我们定义一个列表用来标注蛋白质序列上的 InterPro 功能域或保守位点

interpro_function_annotations = [
    FunctionAnnotation(label="IPR050145", start=1, end=142),  # 1 indexed, inclusive;
    FunctionAnnotation(label="IPR002048", start=4, end=75),
    FunctionAnnotation(label="IPR002048", start=77, end=144),
    FunctionAnnotation(label="IPR011992", start=1, end=143),
    FunctionAnnotation(label="IPR018247", start=17, end=29),
    FunctionAnnotation(label="IPR018247", start=53, end=65),
    FunctionAnnotation(label="IPR018247", start=90, end=102),
    FunctionAnnotation(label="IPR018247", start=126, end=138),
]

接下来我们写一个用于可视化InterPro函数注释的函数

from dna_features_viewer import GraphicFeature, GraphicRecord
from esm.utils.function.interpro import InterPro, InterProEntryType
from matplotlib import colormaps
def visualize_function_annotations(
    annotations: list[FunctionAnnotation],
    sequence_length: int,
    ax: plt.Axes,
    interpro_=InterPro(),
):
    cmap = colormaps["tab10"]
    colors = [cmap(i) for i in range(len(InterProEntryType))]
    type_colors = dict(zip(InterProEntryType, colors))
    features = []
    for annotation in annotations:
        if annotation.label in interpro_.entries:
            entry = interpro_.entries[annotation.label]
            label = entry.name
            entry_type = entry.type
        else:
            label = annotation.label
            entry_type = InterProEntryType.UNKNOWN
        feature = GraphicFeature(
            start=annotation.start - 1,  # one index -> zero index
            end=annotation.end,
            label=label,
            color=type_colors[entry_type],
            strand=None,
        )
        features.append(feature)
    record = GraphicRecord(
        sequence=None, sequence_length=sequence_length, features=features
    )
    record.plot(figure_width=12, plot_sequence=False, ax=ax)

该函数将给定的 InterPro 功能注释列表映射成彩色 GraphicFeature 对象，并借助 dna_features_viewer 在指定 Matplotlib 轴上绘制出整条序列的功能域分布示意图。

打印一下结果：

在使用ESM3 模型时，官方建议您采用关键词注释：这些关键词来自 InterPro 条目描述中的关键词，以及 InterPro2GO 提供的相应 Gene Ontology（GO）术语。例如，针对 InterPro 条目 IPR011992，其关键词包括 “domain pair” 、“hand domain” 、“EF hand” 、“pair” 和 “EF”。关于关键词的具体计算方法，请参阅他们的的预印本论文(https://www.biorxiv.org/content/10.1101/2024.07.01.600583v1.full.pdf)。

在实际操作中，可以利用下述函数根据 InterPro 注释生成关键词注释——每一条 InterPro 注释都对应若干条覆盖同一区间的关键词注释。

from esm.tokenization import InterProQuantizedTokenizer
def get_keywords_from_interpro(
    interpro_annotations,
    interpro2keywords=InterProQuantizedTokenizer().interpro2keywords,
):
    keyword_annotations_list = []
    for interpro_annotation in interpro_annotations:
        keywords = interpro2keywords.get(interpro_annotation.label, [])
        keyword_annotations_list.extend(
            [
                FunctionAnnotation(
                    label=keyword,
                    start=interpro_annotation.start,
                    end=interpro_annotation.end,
                )
                for keyword in keywords
            ]
        )
    return keyword_annotations_list

protein.function_annotations = get_keywords_from_interpro(interpro_function_annotations)
protein.function_annotations

输出：

再来画个图看看：

2.7 sasa溶剂可及表面积

ESMProtein 的最后一条输入通道是“溶剂可及表面积”（SASA）。它为每个氨基酸标示出暴露于溶剂的面积大小。我们可以调用 ProteinChain 的 sasa 函数来计算该指标，该函数内部实际依赖 biotite 库的 sasa 实现。

protein.sasa = protein_chain.sasa()
plt.plot(protein.sasa)

我们还可以将这些SASA值映射到结构的3D可视化上，利用我们拥有这种蛋白质的3D坐标这一事实。

首先，我们定义哪些颜色对应哪些值：

cmap = colormaps["cividis"]
clip_sasa_lower = 10
clip_sasa_upper = 90
def plot_heatmap_legend(cmap, clip_sasa_lower, clip_sasa_upper):
    gradient = np.linspace(0, 1, 256)
    gradient = np.vstack((gradient, gradient))
    _, ax = plt.subplots(figsize=(5, 0.3), dpi=350)
    ax.imshow(gradient, aspect="auto", cmap=cmap)
    ax.text(
        0.1,
        -0.3,
        f"{clip_sasa_lower} or lower",
        va="center",
        ha="right",
        fontsize=7,
        transform=ax.transAxes,
    )
    ax.text(
        0.5,
        -0.3,
        f"{(clip_sasa_lower + clip_sasa_upper) // 2}",
        va="center",
        ha="right",
        fontsize=7,
        transform=ax.transAxes,
    )
    ax.text(
        0.9,
        -0.3,
        f"{clip_sasa_upper} or higher",
        va="center",
        ha="left",
        fontsize=7,
        transform=ax.transAxes,
    )
    ax.set_xticklabels([])
    ax.set_yticklabels([])
    ax.set_xticks([])
    ax.set_yticks([])
    plt.show()
plot_heatmap_legend(cmap, clip_sasa_lower, clip_sasa_upper)

接下来带上结构：

def get_color_strings(sasa, clip_sasa_lower, clip_sasa_upper, cmap):
    transformed_sasa = np.clip(sasa, clip_sasa_lower, clip_sasa_upper)
    transformed_sasa = (transformed_sasa - clip_sasa_lower) / (
        clip_sasa_upper - clip_sasa_lower
    )
    rgbas = (cmap(transformed_sasa) * 255).astype(int)
    return [f"rgb({rgba[0]},{rgba[1]},{rgba[2]})" for rgba in rgbas]
def visualize_sasa_3D_protein(
    protein, clip_sasa_lower=clip_sasa_lower, clip_sasa_upper=clip_sasa_upper, cmap=cmap
):
    pdb_string = protein.to_pdb_string()
    plot_heatmap_legend(cmap, clip_sasa_lower, clip_sasa_upper)
    view = py3Dmol.view(width=400, height=400)
    view.addModel(pdb_string, "pdb")
    for res_pos, res_color in enumerate(
        get_color_strings(protein.sasa, clip_sasa_lower, clip_sasa_upper, cmap)
    ):
        view.setStyle(
            {"chain": "A", "resi": res_pos + 1}, {"cartoon": {"color": res_color}}
        )
    view.zoomTo()
    view.render()
    view.center()
    return view