Protst Esm1b

🚀 ProtST - タンパク質表現学習のための革新的フレームワーク

現在のタンパク質言語モデル（PLM）は主に配列に基づいてタンパク質表現を学習しますが、タンパク質の機能を明示的に獲得することはできません。ProtSTフレームワークは、テキスト記述を用いてタンパク質配列の事前学習と理解を強化し、多様な表現学習ベンチマークで優れた性能を発揮します。

🚀 クイックスタート

現在のタンパク質言語モデル（PLM）は主に配列に基づいてタンパク質表現を学習し、共進化情報をうまく捉えますが、タンパク質表現学習の最終目標であるタンパク質の機能を明示的に獲得することはできません。幸いなことに、多くのタンパク質にはそのテキストによる特性記述があり、そこには様々な機能も記述されています。この事実に基づき、まずProtDescribeデータセットを構築して、タンパク質配列をその機能や他の重要な特性のテキスト記述で拡張します。このデータセットに基づいて、ProtSTフレームワークを提案し、バイオメディカルテキストによってタンパク質配列の事前学習と理解を強化します。事前学習中には、単モーダルマスク予測、マルチモーダル表現アライメント、マルチモーダルマスク予測の3種類のタスクを設計し、異なる粒度でタンパク質の特性情報をPLMに取り入れると同時に、PLMの元の表現能力を維持します。下流タスクでは、ProtSTは教師あり学習とゼロショット予測の両方を可能にします。様々な表現学習ベンチマークで、ProtSTによって誘導されたPLMが従来のものよりも優れていることを検証しました。ゼロショット設定では、ゼロショットタンパク質分類におけるProtSTの有効性を示し、また機能注釈なしで大規模データベースから機能的なタンパク質を検索することも可能です。ソースコードとモデルの重みはhttps://github.com/DeepGraphLearning/ProtSTで入手できます。

✨ 主な機能

• ゼロショット性能の向上：Zero-shot ProtST - ESM - 1bは、少数のサンプルでの分類器を上回る性能を発揮します。

💻 使用例

基本的な使用法

以下のスクリプトは、Gaudiを使用してゼロショット分類タスクでProtSTを実行する方法を示しています。

import logging
import functools
from tqdm import tqdm
import torch
from datasets import load_dataset
from transformers import AutoModel, AutoTokenizer, AutoConfig
+ import habana_frameworks.torch

logger = logging.getLogger(__name__)


def tokenize_protein(example, protein_tokenizer=None, padding=None):
    protein_seqs = example["prot_seq"]

-   protein_inputs = protein_tokenizer(protein_seqs, padding=padding, add_special_tokens=True)
+   protein_inputs = protein_tokenizer(protein_seqs, padding="max_length", truncation=True, add_special_tokens=True, max_length=1024)
    example["protein_input_ids"] = protein_inputs.input_ids
    example["protein_attention_mask"] = protein_inputs.attention_mask
    
    return example


def label_embedding(labels, text_tokenizer, text_model, device):
    # embed label descriptions
    label_feature = []
    with torch.inference_mode():
        for label in labels:
            label_input_ids = text_tokenizer.encode(label, max_length=128,
-                                                   truncation=True, add_special_tokens=False)
+                                                   truncation=True, add_special_tokens=False, padding="max_length")
            label_input_ids = [text_tokenizer.cls_token_id] + label_input_ids
            label_input_ids = torch.tensor(label_input_ids, dtype=torch.long, device=device).unsqueeze(0)
            attention_mask = label_input_ids != text_tokenizer.pad_token_id
            attention_mask = attention_mask.to(device)

            text_outputs = text_model(label_input_ids, attention_mask=attention_mask)

-           label_feature.append(text_outputs["text_feature"])
+           label_feature.append(text_outputs["text_feature"].clone())
    label_feature = torch.cat(label_feature, dim=0)
    label_feature = label_feature / label_feature.norm(dim=-1, keepdim=True)

    return label_feature

def zero_shot_eval(logger, device,                                          
                   test_dataset, target_field, protein_model, logit_scale, label_feature):

    # get prediction and target
    test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False)
    preds, targets = [], []
    with torch.inference_mode():
        for data in tqdm(test_dataloader):
            target = data[target_field]
            targets.append(target)

            protein_input_ids = torch.tensor(data["protein_input_ids"], dtype=torch.long, device=device).unsqueeze(0)
            attention_mask = torch.tensor(data["protein_attention_mask"], dtype=torch.long, device=device).unsqueeze(0)
            
            protein_outputs = protein_model(protein_input_ids, attention_mask=attention_mask)

            protein_feature = protein_outputs["protein_feature"]
            protein_feature = protein_feature / protein_feature.norm(dim=-1, keepdim=True)
            pred = logit_scale * protein_feature @ label_feature.t()
            preds.append(pred)

    preds = torch.cat(preds, dim=0)
    targets = torch.tensor(targets, dtype=torch.long, device=device)
    accuracy = (preds.argmax(dim=-1) == targets).float().mean().item()
    logger.warning("Zero-shot accuracy: %.6f" % accuracy)


if __name__ == "__main__":
    # get datasets
    raw_datasets = load_dataset("mila-intel/ProtST-SubcellularLocalization", cache_dir="~/.cache/huggingface/datasets", split='test') # cache_dir defaults to "~/.cache/huggingface/datasets"

-   device = torch.device("cpu")
+   device = torch.device("hpu")
    
    protst_model = AutoModel.from_pretrained("mila-intel/ProtST-esm1b", trust_remote_code=True, torch_dtype=torch.bfloat16).to(device)
    protein_model = protst_model.protein_model
    text_model = protst_model.text_model
    logit_scale = protst_model.logit_scale
+   from habana_frameworks.torch.hpu import wrap_in_hpu_graph
+   protein_model = wrap_in_hpu_graph(protein_model)
+   text_model = wrap_in_hpu_graph(text_model)
    logit_scale.requires_grad = False
    logit_scale = logit_scale.to(device)
    logit_scale = logit_scale.exp()

    protein_tokenizer = AutoTokenizer.from_pretrained("facebook/esm1b_t33_650M_UR50S")
    text_tokenizer = AutoTokenizer.from_pretrained("microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract")

    func_tokenize_protein = functools.partial(tokenize_protein, protein_tokenizer=protein_tokenizer, padding=False)
    test_dataset = raw_datasets.map(
            func_tokenize_protein, batched=False,
            remove_columns=["prot_seq"],
            desc="Running tokenize_proteins on dataset",
        )

    labels = load_dataset("mila-intel/subloc_template", cache_dir="~/.cache/huggingface/datasets")["train"]["name"]

    text_tokenizer.encode(labels[0], max_length=128, truncation=True, add_special_tokens=False)
    label_feature = label_embedding(labels, text_tokenizer, text_model, device)
    zero_shot_eval(logger, device, test_dataset, "localization",
                   protein_model, logit_scale, label_feature)

高度な使用法

[optimum - intel](https://github.com/huggingface/optimum - intel)の最適化を使用してCPU上でProtSTを実行します。

...
    protst_model = AutoModel.from_pretrained("mila-intel/ProtST-esm1b", trust_remote_code=True, torch_dtype=torch.bfloat16).to(device)
    protein_model = protst_model.protein_model
+   import intel_extension_for_pytorch as ipex
+   from optimum.intel.generation.modeling import jit_trace
+   protein_model = ipex.optimize(protein_model, dtype=torch.bfloat16, inplace=True)
+   protein_model = jit_trace(protein_model, "sequence-classification")
...