蛋白质结构建模

背景

蛋白质是由我们的基因编码的，是生命的基本构建单元，具有复杂性和动态性。它们执行各种功能，例如通过抗体抵抗病毒、为视觉感知光线和为微生物和肌肉提供动力。我们的目标是分析蛋白质结构，以更好地理解它们复杂的组成和功能，从而揭示生命的复杂性，并在生物化学、医学和生物技术等领域取得进展。在这个演示中，我们利用最先进的蛋白质语言模型的向量表示与MyScale数据库进行类似蛋白质搜索和蛋白质活性预测。

如何使用演示？

方法

到目前为止，蛋白质语言模型主要关注于分析单个序列进行推理。传统上，计算生物学依赖于为每个相关序列族拟合单独的模型。我们的演示利用ESMFold (opens new window)模型对1200万个蛋白质序列数据集进行编码，并使用MyScale数据库进行向量搜索，使得执行蛋白质搜索和预测生物活性等任务变得简单。

蛋白质编码器

ESMFold是一种通过为蛋白质中的每个位置分配每个氨基酸的概率来预测蛋白质结构的进化尺度语言模型。向量表示是ESMFold的关键组成部分，它作为蛋白质序列的通用编码。该模型使用这个表示来捕捉序列中每个氨基酸的上下文信息以及不同氨基酸之间的关系。这使得ESMFold能够从单个输入序列生成蛋白质的潜在结构的概率分布。通过在训练过程中利用掩码语言建模目标，ESMFold已经学会生成更准确的预测，从而提高了蛋白质结构预测性能。

使用MyScale进行向量搜索

MyScale结合了SQL和向量数据库的优势，可以在数十亿个向量上进行高速搜索。最新的搜索算法消除了类似示例和困难负例挖掘的挑战。困难负例可以在毫秒级别内找到，分类器训练现在只需在网页上点击几下，为您的AI应用程序节省时间并降低成本。此外，通过混合SQL向量查询实现的改进数据管理显著加快了您的研发过程。

在蛋白质搜索和活性预测中的应用

在这个演示中，我们展示了两个应用：蛋白质搜索和蛋白质活性预测。后者涉及使用ESM的预定义嵌入来预测蛋白质突变的生物影响。这两个应用都利用了MyScale中的向量搜索。

安装依赖

此演示主要使用以下库构建，其中包括：

• transformers：运行ESMFold模型
• clickhouse-connect：数据库客户端
• streamlit：Python Web服务器以运行应用程序

要安装必要的先决条件，请使用以下命令：

python3 -m pip install -r requirement.txt

您可以通过fasta文件在https://github.com/facebookresearch/esm#bulk_fasta (opens new window)下载数据，使用以下命令：

python esm/scripts/extract.py esm2_t33_650M_UR50D examples/data/some_proteins.fasta \
  examples/data/some_proteins_emb_esm2 --repr_layers 0 32 33 --include mean per_tok

使用向量构建数据库

检查数据

让我们来看一下fasta文件的结构。文件P62593.fasta包含蛋白质序列和活性的一行，看起来像这样：

id>protein name>activity>protein sequence
>0|beta-lactamase_P20P|1.581033423 > MSIQHFRVALIPFFAAFCLPVFAHPETLVKVKDAEDQLGARVGYIELDLNSGKILESFRPEERFPMMSTFKVLLCGAVLSRVDAGQEQLGRRIHYSQNDLVEYSPVTEKHLTDGMTVRELCSAAITMSDNTAANLLLTTIGGPKELTAFLHNMGDHVTRLDRWEPELNEAIPNDERDTTMPAAMATTLRKLLTGELLTLASRQQLIDWMEADKVAGPLLRSALPAGWFIADKSGAGERGSRGIIAALGPDGKPSRIVVIYTTGSQATMDERNRQIAEIGASLIKHW

按照我们刚刚分析的数据格式，可以使用pysam轻松获取这些元信息：

from pysam import FastaFile
fasta = "P62593.fasta"
# 读取FASTA文件
sequences_object = FastaFile(fasta)

创建MyScale表

在继续之前，您需要一个有效的凭据来登录我们的数据库引擎。请查看此页面上关于Python客户端的详细指南，了解如何登录。

下面是SQL中表的结构：

CREATE TABLE esm_protein_indexer
( 
    id UInt64,
    activity Float32,
    seq String, 
    representations Array(Float32),
    CONSTRAINT representations CHECK length(vector) = 768
) 

提取特征并填充数据库

ESM模型具有 zero-shot 能力，并且在从蛋白质序列中提取语义特征方面表现出色。它提供了单个蛋白质特征提取和批处理方法。实现代码如下所示。

import torch
import esm
# Load ESM-2 model
model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
batch_converter = alphabet.get_batch_converter()model.eval()  # disables dropout for deterministic results
# Prepare data (first 2 sequences from ESMStructuralSplitDataset superfamily / 4)
data = [
    ("protein1", "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"),
    ("protein2", "KALTARQQEVFDLIRDHISQTGMPPTRAEIAQRLGFRSPNAAEEHLKALARKGVIEIVSGASRGIRLLQEE"),
    ("protein2 with mask","KALTARQQEVFDLIRD<mask>ISQTGMPPTRAEIAQRLGFRSPNAAEEHLKALARKGVIEIVSGASRGIRLLQEE"),
    ("protein3",  "K A <mask> I S Q"),
]
batch_labels, batch_strs, batch_tokens = batch_converter(data)
batch_lens = (batch_tokens != alphabet.padding_idx).sum(1)
# Extract per-residue representations (on CPU)
with torch.no_grad():
    results = model(batch_tokens, repr_layers=[33], return_contacts=True)
    token_representations = results["representations"][33]

我们还可以从FASTA文件中批量提取嵌入。

python scripts/extract.py esm2_t33_650M_UR50D examples/data/some_proteins.fasta \
  examples/data/some_proteins_emb_esm2 --repr_layers 0 32 33 --include mean per_tok

我们已经完成了一个以蛋白质序列特征作为分类器的零-shot学习流程。现在，我认为我们离构建表格已经很接近了。这个拼图中只剩下一个部分：将数据插入MyScale。它看起来像这样：

# 您需要将特征向量转换为Python列表
fields = ['id', 'seq','representations','activity']
# 将向量插入正常的SQL
client.insert("esm_protein_indexer_768", data, column_names=fields)

有关详细的INSERT子句用法，请参考官方文档。

蛋白质搜索

要使用从ESM模型获得的蛋白质嵌入进行搜索，我们提取查询序列的嵌入，并使用SQL在MyScale中定位五个最接近的蛋白质序列的嵌入。然后返回结果。

def esm_search(client, model, sequnce, batch_converter, top_k=5):
    data = [("protein1", sequnce)]
    batch_labels, batch_strs, batch_tokens = batch_converter(data)
    with torch.no_grad():
        results = model(batch_tokens, repr_layers=[33], return_contacts=True)
    token_representations = results["representations"][33]
    token_list = token_representations.tolist()[0][0][0]
    query = f"SELECT activity, distance(representations, {token_list}) as dist "
    query += "FROM default.esm_protein_indexer_768"
    query += f" ORDER BY dist LIMIT {top_k}"
    result = client.query(query).named_results()
    result_temp_coords = [i['coords'] for i in result]
    result_temp_seq = [i['seq'] for i in result]
    return result_temp_coords, result_temp_seq

用于活性预测的KNN

KNN方法通过使用10个最近邻的活性的平均值来预测当前蛋白质的活性。这种方法不需要额外的训练和调整，并且具有相对较高的准确性。

MyScale中的数据包含：

• 突变的β-内酰胺酶序列，其中一个氨基酸残基发生突变（与另一个氨基酸交换）
• 目标值在标题的最后一个字段中，描述突变的缩放效果

函数如下所示：

def knn_search(client, sequence):
    model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
    batch_converter = alphabet.get_batch_converter()
    model.eval()
    data = [("protein1", sequence)]
    batch_labels, batch_strs, batch_tokens = batch_converter(data)
    batch_lens = (batch_tokens != alphabet.padding_idx).sum(1)
    with torch.no_grad():
        results = model(batch_tokens, repr_layers=[33], return_contacts=True)
    token_representations = results["representations"][33]
    token_list = token_representations.tolist()[0][0]
    result = client.query(f"SELECT activity, distance(representations, {token_list}) as dist FROM default.esm_protein_indexer ORDER BY dist LIMIT 10").named_results()
    activity_sum = sum(i['activity'] for i in result)
    avg_activity = activity_sum / len(result)
    return avg_activity

最后

参考文献：

1. 进化尺度建模：https://github.com/facebookresearch/esm (opens new window)
2. KNN：https://towardsdatascience.com/machine-learning-basics-with-the-k-nearest-neighbors-algorithm-6a6e71d01761 (opens new window)