clone from Jiqing's repo

README.md

@@ -1,3 +1,170 @@
 ---
-license: apache-2.0
+library_name: transformers
+tags: []
 ---
+
+# Model Card for Model ID
+
+ProtST for binary localization
+
+## Running script
+```python
+from transformers import AutoModel, AutoTokenizer, HfArgumentParser, TrainingArguments, Trainer
+from transformers.data.data_collator import DataCollatorWithPadding
+from transformers.trainer_pt_utils import get_parameter_names
+from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
+from datasets import load_dataset
+import functools
+import numpy as np
+from sklearn.metrics import accuracy_score, matthews_corrcoef
+import sys
+import torch
+import logging
+import datasets
+import transformers
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def create_optimizer(opt_model, lr_ratio=0.1):
+    head_names = []
+    for n, p in opt_model.named_parameters():
+        if "classifier" in n:
+            head_names.append(n)
+        else:
+            p.requires_grad = False
+    # turn a list of tuple to 2 lists
+    for n, p in opt_model.named_parameters():
+        if n in head_names:
+            assert p.requires_grad
+    backbone_names = []
+    for n, p in opt_model.named_parameters():
+        if n not in head_names and p.requires_grad:
+            backbone_names.append(n)
+    # for weight_decay policy, see 
+    # https://github.com/huggingface/transformers/blob/50573c648ae953dcc1b94d663651f07fb02268f4/src/transformers/trainer.py#L947
+    decay_parameters = get_parameter_names(opt_model, ALL_LAYERNORM_LAYERS) # forbidden layer norm
+    decay_parameters = [name for name in decay_parameters if "bias" not in name]
+    # training_args.learning_rate
+    head_decay_parameters = [name for name in head_names if name in decay_parameters]
+    head_not_decay_parameters = [name for name in head_names if name not in decay_parameters]
+    # training_args.learning_rate * model_config.lr_ratio
+    backbone_decay_parameters = [name for name in backbone_names if name in decay_parameters]
+    backbone_not_decay_parameters = [name for name in backbone_names if name not in decay_parameters]
+    optimizer_grouped_parameters = [
+        {
+            "params": [p for n, p in opt_model.named_parameters() if (n in head_decay_parameters and p.requires_grad)],
+            "weight_decay": training_args.weight_decay,
+            "lr": training_args.learning_rate
+        },
+        {
+            "params": [p for n, p in opt_model.named_parameters() if (n in backbone_decay_parameters and p.requires_grad)],
+            "weight_decay": training_args.weight_decay,
+            "lr": training_args.learning_rate * lr_ratio
+        },
+        {
+            "params": [p for n, p in opt_model.named_parameters() if (n in head_not_decay_parameters and p.requires_grad)],
+            "weight_decay": 0.0,
+            "lr": training_args.learning_rate
+        },
+        {
+            "params": [p for n, p in opt_model.named_parameters() if (n in backbone_not_decay_parameters and p.requires_grad)],
+            "weight_decay": 0.0,
+            "lr": training_args.learning_rate * lr_ratio
+        },
+    ]
+    optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(training_args)
+    optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
+
+    return optimizer
+
+def create_scheduler(training_args, optimizer):
+    from transformers.optimization import get_scheduler
+    return get_scheduler(
+            training_args.lr_scheduler_type,
+            optimizer=optimizer if optimizer is None else optimizer,
+            num_warmup_steps=training_args.get_warmup_steps(training_args.max_steps),
+            num_training_steps=training_args.max_steps,
+        )
+
+def compute_metrics(eval_preds):
+    probs, labels = eval_preds
+    preds = np.argmax(probs, axis=-1)
+    result = {"accuracy": accuracy_score(labels, preds), "mcc": matthews_corrcoef(labels, preds)}
+    return result
+
+def preprocess_logits_for_metrics(logits, labels):
+    return torch.softmax(logits, dim=-1)
+
+
+if __name__ == "__main__":
+    device = torch.device("cpu")
+    raw_dataset = load_dataset("Jiqing/ProtST-BinaryLocalization")
+    model = AutoModel.from_pretrained("Jiqing/protst-esm1b-for-sequential-classification", trust_remote_code=True, torch_dtype=torch.bfloat16).to(device)
+    tokenizer = AutoTokenizer.from_pretrained("facebook/esm1b_t33_650M_UR50S")
+
+    output_dir = "/home/jiqingfe/protst/protst_2/ProtST-HuggingFace/output_dir/ProtSTModel/default/ESM-1b_PubMedBERT-abs/240123_015856"
+    training_args = {'output_dir': output_dir, 'overwrite_output_dir': True, 'do_train': True, 'per_device_train_batch_size': 32, 'gradient_accumulation_steps': 1, \
+                     'learning_rate': 5e-05, 'weight_decay': 0, 'num_train_epochs': 100, 'max_steps': -1, 'lr_scheduler_type': 'constant', 'do_eval': True, \
+                     'evaluation_strategy': 'epoch', 'per_device_eval_batch_size': 32, 'logging_strategy': 'epoch', 'save_strategy': 'epoch', 'save_steps': 820, \
+                     'dataloader_num_workers': 0, 'run_name': 'downstream_esm1b_localization_fix', 'optim': 'adamw_torch', 'resume_from_checkpoint': False, \
+                     'label_names': ['labels'], 'load_best_model_at_end': True, 'metric_for_best_model': 'accuracy', 'bf16': True, "save_total_limit": 3}
+    training_args = HfArgumentParser(TrainingArguments).parse_dict(training_args, allow_extra_keys=False)[0]
+
+    def tokenize_protein(example, tokenizer=None):
+        protein_seq = example["prot_seq"]
+        protein_seq_str = tokenizer(protein_seq, add_special_tokens=True)
+        example["input_ids"] = protein_seq_str["input_ids"]
+        example["attention_mask"] = protein_seq_str["attention_mask"]
+        example["labels"] = example["localization"]
+
+        return example
+
+    func_tokenize_protein = functools.partial(tokenize_protein, tokenizer=tokenizer)
+
+    for split in ["train", "validation", "test"]:
+        raw_dataset[split] = raw_dataset[split].map(func_tokenize_protein, batched=False, remove_columns=["Unnamed: 0", "prot_seq", "localization"])
+
+    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+
+    transformers.utils.logging.set_verbosity_info()
+    log_level = training_args.get_process_log_level()
+    logger.setLevel(log_level)
+
+    optimizer = create_optimizer(model)
+    scheduler = create_scheduler(training_args, optimizer)
+
+    # build trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=raw_dataset["train"],
+        eval_dataset=raw_dataset["validation"],
+        data_collator=data_collator,
+        optimizers=(optimizer, scheduler),
+        compute_metrics=compute_metrics,
+        preprocess_logits_for_metrics=preprocess_logits_for_metrics,
+    )
+
+    train_result = trainer.train()
+
+    trainer.save_model()
+    # Saves the tokenizer too for easy upload
+    tokenizer.save_pretrained(training_args.output_dir)
+
+    metrics = train_result.metrics
+    metrics["train_samples"] = len(raw_dataset["train"])
+
+    trainer.log_metrics("train", metrics)
+    trainer.save_metrics("train", metrics)
+    trainer.save_state()
+
+    metric = trainer.evaluate(raw_dataset["test"], metric_key_prefix="test")
+    print("test metric: ", metric)
+
+    metric = trainer.evaluate(raw_dataset["validation"], metric_key_prefix="valid")
+    print("valid metric: ", metric)
+```
+
+
+

config.json

@@ -0,0 +1,37 @@
+{
+  "_name_or_path": "Jiqing/protst-esm1b-for-sequential-classification",
+  "architectures": [
+    "ProtSTForProteinPropertyPrediction"
+  ],
+  "auto_map": {
+    "AutoConfig": "Jiqing/protst-esm1b-for-sequential-classification--configuration_protst.ProtSTConfig",
+    "AutoModel": "Jiqing/protst-esm1b-for-sequential-classification--modeling_protst.ProtSTForProteinPropertyPrediction"
+  },
+  "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,
+    "cls_token_id": 0,
+    "emb_layer_norm_before": true,
+    "eos_token_id": 2,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.0,
+    "hidden_size": 1280,
+    "intermediate_size": 5120,
+    "layer_norm_eps": 1e-05,
+    "mask_token_id": 32,
+    "model_type": "esm",
+    "num_attention_heads": 20,
+    "num_hidden_layers": 33,
+    "pad_token_id": 1,
+    "token_dropout": true,
+    "torch_dtype": "float32",
+    "vocab_size": 33
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.38.0.dev0"
+}

configuration_protst.py

@@ -0,0 +1,42 @@
+from transformers import PretrainedConfig
+from transformers.utils import logging
+from transformers.models.esm import EsmConfig
+
+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`].
+    ```"""
+
+    model_type = "protst"
+
+    def __init__(
+        self,
+        protein_config=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if protein_config is None:
+            protein_config = {}
+            logger.info("`protein_config` is `None`. Initializing the `ProtSTProteinConfig` with default values.")
+
+        self.protein_config = EsmConfig(**protein_config)
+
+    @classmethod
+    def from_protein_text_configs(
+        cls, protein_config: EsmConfig, **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(), **kwargs)

model.safetensors

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

modeling_protst.py

@@ -0,0 +1,213 @@
+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.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)
+
+        return EsmProteinRepresentationOutput(
+            protein_feature=protein_feature, residue_feature=residue_feature
+        )
+
+
+class ProtSTPreTrainedModel(PreTrainedModel):
+    config_class = ProtSTConfig
+
+
+class ProtSTForProteinPropertyPrediction(ProtSTPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.config = config
+        self.protein_model = EsmForProteinRepresentation(config.protein_config)
+        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)