Prostt5

🚀 ProstT5モデルカード

ProstT5は、タンパク質配列と構造の間の翻訳を行うことができるタンパク質言語モデル（pLM）です。このモデルは、タンパク質の構造と配列の関係を理解し、相互変換する能力を持っています。

📚 詳細ドキュメント

モデルの説明

ProstT5（Protein structure - sequence T5）は、ProtT5 - XL - U50をベースにしています。このT5モデルは、何十億ものタンパク質配列に対してスパン破損を適用してタンパク質配列をエンコードするように学習されています。

ProstT5は、AlphaFoldDBからの高品質な3D構造予測を持つ1700万個のタンパク質を使用して、タンパク質配列と構造の間の翻訳に関してProtT5 - XL - U50を微調整します。

タンパク質構造は、Foldseekによって導入された3Diトークンを使用して3Dから1Dに変換されます。最初に、ProstT5は、3Diおよびアミノ酸（AA）配列に適用される元のスパンノイズ除去の目的を継続することで、新しく導入された3Diトークンを表現するように学習します。2番目のステップでのみ、ProstT5は2つのモダリティ間の翻訳について学習されます。

翻訳の方向は、2つの特殊トークン（3DiからAAへの翻訳には"<fold2AA>"、AAから3Diへの翻訳には“<AA2fold>”）によって示されます。AAトークンとの衝突を避けるために、3Diトークンは小文字に変換されます（それ以外のアルファベットは同じです）。

• 開発者: Michael Heinzinger (GitHub @mheinzinger; Twitter @HeinzingerM)
• モデルの種類: エンコーダ - デコーダ（T5）
• 言語 (NLP): タンパク質配列と構造
• ライセンス: MIT
• 微調整元のモデル: ProtT5 - XL - U50

✨ 主な機能

1. このモデルは、従来の特徴抽出に使用することができます。この目的のために、半精度（fp16）でバッチ処理を行うエンコーダのみを使用することをおすすめします。例（現在は元のProtT5 - XL - U50のみですが、リポジトリリンクを置き換えて接頭辞を追加することで動作します）: スクリプト と [colab](https://colab.research.google.com/drive/1h7F5v5xkE_ly - 1bTQSu - 1xaLtTP2TnLF?usp = sharing)

元のProtT5 - XL - U50はAA配列のみを埋め込むことができましたが、ProstT5は現在、3Diトークンで表される3D構造も埋め込むことができます。3Diトークンは、Foldseekを使用して3D構造から導出することも、ProstT5によってAA配列から予測することもできます。 3. "Folding": 配列（AA）から構造（3Di）への翻訳。結果として得られる3Di文字列は、Foldseekと一緒に使用して、3D構造を明示的に計算することなく遠縁相同性検出を行うことができます。 4. "Inverse Folding": 構造（3Di）から配列（AA）への翻訳。

🚀 クイックスタート

特徴抽出

from transformers import T5Tokenizer, T5EncoderModel
import torch
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# Load the tokenizer
tokenizer = T5Tokenizer.from_pretrained('Rostlab/ProstT5', do_lower_case=False).to(device)

# Load the model
model = T5EncoderModel.from_pretrained("Rostlab/ProstT5").to(device)

# only GPUs support half-precision currently; if you want to run on CPU use full-precision (not recommended, much slower)
model.full() if device=='cpu' else model.half()

# prepare your protein sequences/structures as a list. Amino acid sequences are expected to be upper-case ("PRTEINO" below) while 3Di-sequences need to be lower-case ("strctr" below).
sequence_examples = ["PRTEINO", "strct"]

# replace all rare/ambiguous amino acids by X (3Di sequences does not have those) and introduce white-space between all sequences (AAs and 3Di)
sequence_examples = [" ".join(list(re.sub(r"[UZOB]", "X", sequence))) for sequence in sequence_examples]

# add pre-fixes accordingly (this already expects 3Di-sequences to be lower-case)
# if you go from AAs to 3Di (or if you want to embed AAs), you need to prepend "<AA2fold>"
# if you go from 3Di to AAs (or if you want to embed 3Di), you need to prepend "<fold2AA>"
sequence_examples = [ "<AA2fold>" + " " + s if s.isupper() else "<fold2AA>" + " " + s
                      for s in sequence_examples
                    ]

# tokenize sequences and pad up to the longest sequence in the batch
ids = tokenizer.batch_encode_plus(sequences_example, add_special_tokens=True, padding="longest",return_tensors='pt').to(device))

# generate embeddings
with torch.no_grad():
    embedding_rpr = model(
              ids.input_ids, 
              attention_mask=ids.attention_mask
              )

# extract residue embeddings for the first ([0,:]) sequence in the batch and remove padded & special tokens, incl. prefix ([0,1:8]) 
emb_0 = embedding_repr.last_hidden_state[0,1:8] # shape (7 x 1024)
# same for the second ([1,:]) sequence but taking into account different sequence lengths ([1,:6])
emb_1 = embedding_repr.last_hidden_state[1,1:6] # shape (5 x 1024)

# if you want to derive a single representation (per-protein embedding) for the whole protein
emb_0_per_protein = emb_0.mean(dim=0) # shape (1024)

翻訳 ("folding", すなわちAAから3Di)

from transformers import T5Tokenizer, AutoModelForSeq2SeqLM
import torch
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# Load the tokenizer
tokenizer = T5Tokenizer.from_pretrained('Rostlab/ProstT5', do_lower_case=False).to(device)

# Load the model
model = AutoModelForSeq2SeqLM.from_pretrained("Rostlab/ProstT5").to(device)

# only GPUs support half-precision currently; if you want to run on CPU use full-precision (not recommended, much slower)
model.full() if device=='cpu' else model.half()

# prepare your protein sequences/structures as a list.
# Amino acid sequences are expected to be upper-case ("PRTEINO" below)
# while 3Di-sequences need to be lower-case.
sequence_examples = ["PRTEINO", "SEQWENCE"]
min_len = min([ len(s) for s in folding_example])
max_len = max([ len(s) for s in folding_example])

# replace all rare/ambiguous amino acids by X (3Di sequences does not have those) and introduce white-space between all sequences (AAs and 3Di)
sequence_examples = [" ".join(list(re.sub(r"[UZOB]", "X", sequence))) for sequence in sequence_examples]

# add pre-fixes accordingly. For the translation from AAs to 3Di, you need to prepend "<AA2fold>"
sequence_examples = [ "<AA2fold>" + " " + s for s in sequence_examples]

# tokenize sequences and pad up to the longest sequence in the batch
ids = tokenizer.batch_encode_plus(sequences_example,
                                  add_special_tokens=True,
                                  padding="longest",
                                  return_tensors='pt').to(device))

# Generation configuration for "folding" (AA-->3Di)
gen_kwargs_aa2fold = {
                  "do_sample": True,
                  "num_beams": 3, 
                  "top_p" : 0.95, 
                  "temperature" : 1.2, 
                  "top_k" : 6,
                  "repetition_penalty" : 1.2,
}

# translate from AA to 3Di (AA-->3Di)
with torch.no_grad():
  translations = model.generate( 
              ids.input_ids, 
              attention_mask=ids.attention_mask, 
              max_length=max_len, # max length of generated text
              min_length=min_len, # minimum length of the generated text
              early_stopping=True, # stop early if end-of-text token is generated
              num_return_sequences=1, # return only a single sequence
              **gen_kwargs_aa2fold
  )
# Decode and remove white-spaces between tokens
decoded_translations = tokenizer.batch_decode( translations, skip_special_tokens=True )
structure_sequences = [ "".join(ts.split(" ")) for ts in decoded_translations ] # predicted 3Di strings

# Now we can use the same model and invert the translation logic
# to generate an amino acid sequence from the predicted 3Di-sequence (3Di-->AA)

# add pre-fixes accordingly. For the translation from 3Di to AA (3Di-->AA), you need to prepend "<fold2AA>"
sequence_examples_backtranslation = [ "<fold2AA>" + " " + s for s in decoded_translations]

# tokenize sequences and pad up to the longest sequence in the batch
ids_backtranslation = tokenizer.batch_encode_plus(sequence_examples_backtranslation,
                                  add_special_tokens=True,
                                  padding="longest",
                                  return_tensors='pt').to(device))

# Example generation configuration for "inverse folding" (3Di-->AA)
gen_kwargs_fold2AA = {
            "do_sample": True,
            "top_p" : 0.90,
            "temperature" : 1.1,
            "top_k" : 6,
            "repetition_penalty" : 1.2,
}

# translate from 3Di to AA (3Di-->AA)
with torch.no_grad():
  backtranslations = model.generate( 
              ids_backtranslation.input_ids, 
              attention_mask=ids_backtranslation.attention_mask, 
              max_length=max_len, # max length of generated text
              min_length=min_len, # minimum length of the generated text
              early_stopping=True, # stop early if end-of-text token is generated
              num_return_sequences=1, # return only a single sequence
              **gen_kwargs_fold2AA
  )
# Decode and remove white-spaces between tokens
decoded_backtranslations = tokenizer.batch_decode( backtranslations, skip_special_tokens=True )
aminoAcid_sequences = [ "".join(ts.split(" ")) for ts in decoded_backtranslations ] # predicted amino acid strings

🔧 技術詳細

学習データ

事前学習データ (1700万個のタンパク質の3Di + AA配列)

学習手順

事前学習の最初のフェーズでは、このスクリプトを使用して、3DiおよびAA配列を用いたスパンベースのノイズ除去を続けます。

事前学習の2番目のフェーズ（実際の3DiからAA配列への翻訳およびその逆）では、このスクリプトを使用しました。

学習ハイパーパラメータ

• 学習方式: DeepSpeed（stage - 2）、勾配累積ステップ（5ステップ）、混合半精度（bf16）およびPyTorch2.0のtorchInductorコンパイラを使用しました

速度

Pro(s)tT5エンコーダからヒトプロテオームの埋め込みを生成するには、48GBのvRAMを持つ単一のRTX A6000 GPUでバッチ処理と半精度（fp16）を使用すると、約35分（分）またはタンパク質あたり0.1秒（秒）かかります。

翻訳は比較的遅く（それぞれ平均長135と406で0.6 - 2.5秒/タンパク質）、左から右にトークンごとに生成する必要があるデコードプロセスの逐次的な性質によるものです。

さらなる最適化を行わずに、半精度でバッチ処理のみを使用しました。

📄 ライセンス

このモデルはMITライセンスの下で提供されています。