Saprot 650M AF2

🚀 SaProt

SaProtは、タンパク質関連のタスクで使用できるモデルです。構造入力を用いることで、多くのタスクにおいてESM2を上回る性能を発揮します。本READMEでは、SaProtの使用方法を2通り提供しています。

🚀 クイックスタート

SaProtを使用する方法は2通りあり、huggingfaceクラスを通じて使用する方法と、esm githubと同じ方法で使用する方法があります。ユーザーはいずれかを選択して使用できます。

✨ 主な機能

• Huggingfaceを通じたモデルの読み込みと使用が可能です。
• esmと同じ方法での使用もサポートしています。
• タンパク質配列の変異効果の予測ができます。
• タンパク質の埋め込みを取得することができます。

📦 インストール

インストールに関する具体的な手順は原ドキュメントに記載されていないため、このセクションをスキップします。

💻 使用例

基本的な使用法

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)

#################### Example ####################
device = "cuda"
model.to(device)

seq = "M#EvVpQpL#VyQdYaKv" # Here "#" represents lower plDDT regions (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", # Note this is the directory path of SaProt, not the ".pt" file
    "load_pretrained": True,
}
model = SaprotFoldseekMutationModel(**config)
tokenizer = model.tokenizer

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

seq = "M#EvVpQpL#VyQdYaKv" # Here "#" represents lower plDDT regions (plddt < 70)

# Predict the effect of mutating the 3rd amino acid to A
mut_info = "V3A"
mut_value = model.predict_mut(seq, mut_info)
print(mut_value)

# Predict mutational effect of combinatorial mutations, e.g. mutating the 3rd amino acid to A and the 4th amino acid to M
mut_info = "V3A:Q4M"
mut_value = model.predict_mut(seq, mut_info)
print(mut_value)

# Predict all effects of mutations at 3rd position
mut_pos = 3
mut_dict = model.predict_pos_mut(seq, mut_pos)
print(mut_dict)

# Predict probabilities of all amino acids at 3rd position
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", # Note this is the directory path of SaProt, not the ".pt" file
    "load_pretrained": True,
}

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

device = "cuda"
model.to(device)

seq = "M#EvVpQpL#VyQdYaKv" # Here "#" represents lower plDDT regions (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ライセンスの下で提供されています。

⚠️ 重要提示

SaProtは構造的（SAトークン）入力を必要とし、最適な性能を発揮します。アミノ酸配列のみのモードも動作しますが、微調整が必要です。凍結された埋め込みはSAにのみ機能し、アミノ酸配列には機能しません。構造入力を使用すると、SaProtはほとんどのタスクでESM2を上回ります。