Saprot 650M AF2

🚀 SaProt模型使用說明

SaProt是一個在蛋白質相關任務中表現出色的模型，它需要結構（SA令牌）輸入以達到最佳性能。僅使用氨基酸序列模式雖然可行，但必須進行微調，因為凍結的嵌入僅適用於SA，而不適用於氨基酸序列。在有結構輸入的情況下，SaProt在大多數任務中超越了ESM2。本項目提供了兩種使用SaProt的方式，用戶可按需選擇。

🚀 快速開始

我們提供了兩種使用SaProt的方式，包括通過Huggingface類調用，以及採用與 esm github 相同的方式。用戶可以任選其一使用。

✨ 主要特性

• 支持通過Huggingface類調用模型。
• 支持以與esm相同的方式加載模型。
• 提供預測蛋白質序列突變效應的功能。
• 可生成蛋白質嵌入。

📦 安裝指南

文檔未提及具體安裝步驟，故跳過此章節。

💻 使用示例

基礎用法

Huggingface模型調用

以下代碼展示瞭如何加載Huggingface模型：

from transformers import EsmTokenizer, EsmForMaskedLM

model_path = "/your/path/to/SaProt_650M_AF2"
tokenizer = EsmTokenizer.from_pretrained(model_path)
model = EsmForMaskedLM.from_pretrained(model_path)

#################### 示例 ####################
device = "cuda"
model.to(device)

seq = "M#EvVpQpL#VyQdYaKv" # 這裡 "#" 代表較低的pLDDT區域 (plddt < 70)
tokens = tokenizer.tokenize(seq)
print(tokens)

inputs = tokenizer(seq, return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}

outputs = model(**inputs)
print(outputs.logits.shape)

"""
['M#', 'Ev', 'Vp', 'Qp', 'L#', 'Vy', 'Qd', 'Ya', 'Kv']
torch.Size([1, 11, 446])
"""

esm模型調用

esm版本的模型也存儲在同一文件夾中，名為 SaProt_650M_AF2.pt。我們提供了一個函數來加載該模型：

from utils.esm_loader import load_esm_saprot

model_path = "/your/path/to/SaProt_650M_AF2.pt"
model, alphabet = load_esm_saprot(model_path)

高級用法

預測突變效應

我們提供了一個函數來預測蛋白質序列的突變效應。以下示例展示瞭如何預測特定位置的突變效應。如果使用AF2結構，強烈建議添加pLDDT掩碼（見下文）：

from model.saprot.saprot_foldseek_mutation_model import SaprotFoldseekMutationModel

config = {
    "foldseek_path": None,
    "config_path": "/your/path/to/SaProt_650M_AF2", # 注意這是SaProt的目錄路徑，而不是 ".pt" 文件
    "load_pretrained": True,
}
model = SaprotFoldseekMutationModel(**config)
tokenizer = model.tokenizer

device = "cuda"
model.eval()
model.to(device)

seq = "M#EvVpQpL#VyQdYaKv" # 這裡 "#" 代表較低的pLDDT區域 (plddt < 70)

# 預測將第3個氨基酸突變為A的效應
mut_info = "V3A"
mut_value = model.predict_mut(seq, mut_info)
print(mut_value)

# 預測組合突變的效應，例如將第3個氨基酸突變為A，第4個氨基酸突變為M
mut_info = "V3A:Q4M"
mut_value = model.predict_mut(seq, mut_info)
print(mut_value)

# 預測第3個位置所有突變的效應
mut_pos = 3
mut_dict = model.predict_pos_mut(seq, mut_pos)
print(mut_dict)

# 預測第3個位置所有氨基酸的概率
mut_pos = 3
mut_dict = model.predict_pos_prob(seq, mut_pos)
print(mut_dict)

獲取蛋白質嵌入

如果想生成蛋白質嵌入，可以參考以下代碼。嵌入是最後一層隱藏狀態的平均值：

from model.saprot.base import SaprotBaseModel
from transformers import EsmTokenizer

config = {
    "task": "base",
    "config_path": "/your/path/to/SaProt_650M_AF2", # 注意這是SaProt的目錄路徑，而不是 ".pt" 文件
    "load_pretrained": True,
}

model = SaprotBaseModel(**config)
tokenizer = EsmTokenizer.from_pretrained(config["config_path"])

device = "cuda"
model.to(device)

seq = "M#EvVpQpL#VyQdYaKv" # 這裡 "#" 代表較低的pLDDT區域 (plddt < 70)
tokens = tokenizer.tokenize(seq)
print(tokens)

inputs = tokenizer(seq, return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}

embeddings = model.get_hidden_states(inputs, reduction="mean")
print(embeddings[0].shape)

📚 詳細文檔

文檔未提供詳細說明，故跳過此章節。

🔧 技術細節

文檔未提供技術實現細節，故跳過此章節。

📄 許可證

本項目採用MIT許可證。