init from Jiqing's repo

README.md

@@ -1,3 +1,108 @@
----
-license: apache-2.0
----
+
+## Abstract
+Current protein language models (PLMs) learn protein representations mainly based on their sequences, thereby well capturing co-evolutionary information, but they are unable to explicitly acquire protein functions, which is the end goal of protein representation learning. Fortunately, for many proteins, their textual property descriptions are available, where their various functions are also described. Motivated by this fact, we first build the ProtDescribe dataset to augment protein sequences with text descriptions of their functions and other important properties. Based on this dataset, we propose the [ProtST framework](https://arxiv.org/abs/2301.12040) to enhance Protein Sequence pre-training and understanding by biomedical Texts. During pre-training, we design three types of tasks, i.e., unimodal mask prediction, multimodal representation alignment and multimodal mask prediction, to enhance a PLM with protein property information with different granularities and, at the same time, preserve the PLM’s original representation power. On downstream tasks, ProtST enables both supervised learning and zeroshot prediction. We verify the superiority of ProtST-induced PLMs over previous ones on diverse representation learning benchmarks. Under the zero-shot setting, we show the effectiveness of ProtST on zero-shot protein classification, and ProtST also enables functional protein retrieval from a large-scale database without any function annotation. Source code and model weights are available at [https://github.com/DeepGraphLearning/ProtST](https://github.com/DeepGraphLearning/ProtST).
+
+![image/png](https://cdn-uploads.huggingface.co/production/uploads/62f0a673f0d40f6aae296b4a/o4F5-Cm-gGdHPpX5rPVKx.png)
+
+## Example
+This example shows how to use ProtST on zero-shot classification task.
+```python
+import logging
+import functools
+from tqdm import tqdm
+import torch
+from datasets import load_dataset
+from transformers import AutoModel, AutoTokenizer, AutoConfig
+
+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)
+    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)
+            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 = 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("Jiqing/ProtST-SubcellularLocalization", cache_dir="~/.cache/huggingface/datasets", split='test') # cache_dir defaults to "~/.cache/huggingface/datasets"
+    
+    #device = torch.device("cuda:0")
+    device = torch.device("cpu")
+    
+    protst_model = AutoModel.from_pretrained("Jiqing/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
+    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("Jiqing/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)
+```

config.json

@@ -0,0 +1,63 @@
+{
+  "architectures": [
+    "ProtSTModel"
+  ],
+  "auto_map": {
+    "AutoModel": "modeling_protst.ProtSTModel",
+    "AutoConfig": "configuration_protst.ProtSTConfig"
+  },
+  "model_type": "protst",
+  "protein_config": {
+    "_name_or_path": "/tmp/facebook/esm1b_t33_650M_UR50S",
+    "architectures": [
+      "EsmForMaskedLM"
+    ],
+    "attention_probs_dropout_prob": 0.0,
+    "classifier_dropout": null,
+    "emb_layer_norm_before": true,
+    "esmfold_config": null,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.0,
+    "hidden_size": 1280,
+    "initializer_range": 0.02,
+    "intermediate_size": 5120,
+    "is_folding_model": false,
+    "layer_norm_eps": 1e-05,
+    "mask_token_id": 32,
+    "max_position_embeddings": 1026,
+    "model_type": "esm",
+    "num_attention_heads": 20,
+    "num_hidden_layers": 33,
+    "cls_token_id": 0,
+    "pad_token_id": 1,
+    "eos_token_id": 2,
+    "position_embedding_type": "absolute",
+    "token_dropout": true,
+    "torch_dtype": "float32",
+    "use_cache": true,
+    "vocab_list": null,
+    "vocab_size": 33
+  },
+  "text_config": {
+    "architectures": [
+      "BertForMaskedLM"
+    ],
+    "model_type": "bert",
+    "attention_probs_dropout_prob": 0.1,
+    "hidden_act": "gelu",
+    "pad_token_id": 0,
+    "cls_token_id": 2,
+    "sep_token_id": 3,
+    "hidden_dropout_prob": 0.1,
+    "hidden_size": 768,
+    "initializer_range": 0.02,
+    "intermediate_size": 3072,
+    "max_position_embeddings": 512,
+    "num_attention_heads": 12,
+    "num_hidden_layers": 12,
+    "type_vocab_size": 2,
+    "vocab_size": 30522
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.37.0.dev0"
+}

configuration_protst.py

@@ -0,0 +1,53 @@
+from transformers import PretrainedConfig
+from transformers.utils import logging
+from transformers.models.esm import EsmConfig
+from transformers.models.bert import BertConfig
+
+logger = logging.get_logger(__name__)
+
+
+class ProtSTConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`ProtSTModel`].
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        protein_config (`dict`, *optional*):
+            Dictionary of configuration options used to initialize [`EsmForProteinRepresentation`].
+        text_config (`dict`, *optional*):
+            Dictionary of configuration options used to initialize [`BertForPubMed`].
+    ```"""
+
+    model_type = "protst"
+
+    def __init__(
+        self,
+        protein_config=None,
+        text_config=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if protein_config is None:
+            protein_config = {}
+            logger.info("`protein_config` is `None`. Initializing the `ProtSTTextConfig` with default values.")
+
+        if text_config is None:
+            text_config = {}
+            logger.info("`text_config` is `None`. Initializing the `ProtSTVisionConfig` with default values.")
+
+        self.protein_config = EsmConfig(**protein_config)
+        self.text_config = BertConfig(**text_config)
+
+    @classmethod
+    def from_protein_text_configs(
+        cls, protein_config: EsmConfig, text_config: BertConfig, **kwargs
+    ):
+        r"""
+        Instantiate a [`ProtSTConfig`] (or a derived class) from ProtST text model configuration. Returns:
+            [`ProtSTConfig`]: An instance of a configuration object
+        """
+
+        return cls(protein_config=protein_config.to_dict(), text_config=text_config.to_dict(), **kwargs)

modeling_protst.py

@@ -0,0 +1,285 @@
+import math
+import torch
+import torch.nn as nn
+from typing import Optional, Tuple, Union
+from dataclasses import dataclass
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+from transformers.models.esm import EsmPreTrainedModel, EsmModel
+from transformers.models.bert import BertPreTrainedModel, BertModel
+from .configuration_protst import ProtSTConfig
+
+
+@dataclass
+class EsmProteinRepresentationOutput(ModelOutput):
+
+    protein_feature: torch.FloatTensor = None
+    residue_feature: torch.FloatTensor = None
+
+
+@dataclass
+class BertTextRepresentationOutput(ModelOutput):
+
+    text_feature: torch.FloatTensor = None
+    word_feature: torch.FloatTensor = None
+
+
+@dataclass
+class ProtSTClassificationOutput(ModelOutput):
+
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+
+class ProtSTHead(nn.Module):
+    def __init__(self, config, out_dim=512):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.out_proj = nn.Linear(config.hidden_size, out_dim)
+
+    def forward(self, x):
+        x = self.dense(x)
+        x = nn.functional.relu(x)
+        x = self.out_proj(x)
+        return x
+
+
+class BertForPubMed(BertPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.pad_token_id = config.pad_token_id
+        self.cls_token_id = config.cls_token_id
+        self.sep_token_id = config.sep_token_id
+
+        self.bert = BertModel(config, add_pooling_layer=False)
+        self.text_mlp = ProtSTHead(config)
+        self.word_mlp = ProtSTHead(config)
+
+        self.post_init() # NOTE
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor], ModelOutput]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        word_feature = outputs.last_hidden_state
+        is_special = (input_ids == self.cls_token_id) | (input_ids == self.sep_token_id) | (input_ids == self.pad_token_id)
+        special_mask = (~is_special).to(torch.int64).unsqueeze(-1)
+        pooled_feature = ((word_feature * special_mask).sum(1) / (special_mask.sum(1) + 1.0e-6)).to(word_feature.dtype)
+        pooled_feature = self.text_mlp(pooled_feature)
+        word_feature = self.word_mlp(word_feature)
+
+        if not return_dict:
+            return (pooled_feature, word_feature)
+
+        return BertTextRepresentationOutput(text_feature=pooled_feature, word_feature=word_feature)
+        
+
+
+
+class EsmForProteinRepresentation(EsmPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.cls_token_id = config.cls_token_id
+        self.pad_token_id = config.pad_token_id
+        self.eos_token_id = config.eos_token_id
+
+        self.esm = EsmModel(config, add_pooling_layer=False)
+        self.protein_mlp = ProtSTHead(config)
+        self.residue_mlp = ProtSTHead(config)
+
+        self.post_init() # NOTE
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, EsmProteinRepresentationOutput]:
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.esm(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        residue_feature = outputs.last_hidden_state  # [batch_size, seq_len, hidden_dim]
+
+        # mean readout
+        is_special = (
+            (input_ids == self.cls_token_id) | (input_ids == self.eos_token_id) | (input_ids == self.pad_token_id)
+        )
+        special_mask = (~is_special).to(torch.int64).unsqueeze(-1)
+        protein_feature = ((residue_feature * special_mask).sum(1) / (special_mask.sum(1) + 1.0e-6)).to(residue_feature.dtype)
+
+        # For ProtST pretrain and zero-shot
+        protein_feature = self.protein_mlp(protein_feature)
+        residue_feature = self.residue_mlp(residue_feature)
+
+        
+        return EsmProteinRepresentationOutput(
+            protein_feature=protein_feature, residue_feature=residue_feature
+        )
+
+
+class ProtSTPreTrainedModel(PreTrainedModel):
+    config_class = ProtSTConfig
+
+    def _compute_protein_feature(self, 
+        protein_input_ids, protein_attention_mask, protein_position_ids, 
+        output_attentions, output_hidden_states
+    ):
+
+        protein_outputs = self.protein_model(
+            protein_input_ids,
+            attention_mask=protein_attention_mask,
+            position_ids=protein_position_ids,
+            head_mask=None,
+            inputs_embeds=None,
+            encoder_hidden_states=None,
+            encoder_attention_mask=None,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=None,
+        )
+        
+        return protein_outputs
+
+    def _compute_text_feature(self, 
+        text_input_ids, text_attention_mask, text_position_ids,
+        output_attentions, output_hidden_states
+    ):
+        text_outputs = self.text_model(
+            text_input_ids,
+            attention_mask=text_attention_mask,
+            position_ids=text_position_ids,
+            head_mask=None,
+            inputs_embeds=None,
+            encoder_hidden_states=None,
+            encoder_attention_mask=None,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=None,
+        )
+
+        return text_outputs
+
+
+class ProtSTModel(ProtSTPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.config = config
+        self.protein_model = EsmForProteinRepresentation(config.protein_config)
+        self.text_model = BertForPubMed(config.text_config)
+        self.logit_scale = nn.Parameter(torch.ones([]) * math.log(1 / 0.07))
+
+        self.post_init() # NOTE
+
+    def forward(self, 
+        protein_input_ids: Optional[torch.LongTensor] = None,
+        text_input_ids: Optional[torch.LongTensor] = None,
+        protein_attention_mask: Optional[torch.Tensor] = None,
+        text_attention_mask: Optional[torch.Tensor] = None,
+        protein_position_ids: Optional[torch.LongTensor] = None,
+        text_position_ids: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+    ):
+        # Not implement yet
+        return None
+
+
+class ProtSTForProteinPropertyPrediction(ProtSTPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.config = config
+        self.protein_model = EsmForProteinRepresentation(config.protein_config)
+        self.logit_scale = nn.Parameter(torch.ones([]) * math.log(1 / 0.07))
+        self.classifier = ProtSTHead(config.protein_config, out_dim=config.num_labels)
+
+        self.post_init() # NOTE
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, ProtSTClassificationOutput]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the protein classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
+        Returns:
+        Examples:
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.protein_model(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        logits = self.classifier(outputs.protein_feature) # [bsz, xxx] -> [bsz, num_labels]
+
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+
+            labels = labels.to(logits.device)
+            loss = loss_fct(logits.view(-1, logits.shape[-1]), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,)
+            return ((loss,) + output) if loss is not None else output
+
+        return ProtSTClassificationOutput(loss=loss, logits=logits)

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c59f77e12992626701f6bdfb732b5b9171f753fda86df7f68aa2135ebd421868
+size 135