jiaxie/DeepProtT5-Fold

First we define a class T5ClassificationModel:

from transformers import (
    T5Config,
    T5EncoderModel,
    T5Tokenizer,
    PreTrainedModel,
    TrainingArguments,
    Trainer,
    DataCollatorWithPadding,
)
class T5ClassificationModel(PreTrainedModel):
    config_class = T5Config

    def __init__(self, config, d_model=None, num_classes=1195):
        super().__init__(config)
        self.num_classes = num_classes

        self.encoder = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_uniref50")

        hidden_dim = d_model if d_model is not None else config.d_model
        self.classification_head = nn.Linear(hidden_dim, num_classes)

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        labels=None,
        **kwargs
    ):
        encoder_outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        hidden_states = encoder_outputs.last_hidden_state

        mask = attention_mask.unsqueeze(-1)
        pooled_output = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1)
        logits = self.classification_head(pooled_output)  # [batch_size, num_classes]

        loss = None
        if labels is not None:
            labels = labels.to(torch.bfloat16)
            loss = nn.CrossEntropyLoss()(logits, labels)

        return {
            "loss": loss,
            "logits": logits
        }

Then we load our pretrained model

tokenizer = T5Tokenizer.from_pretrained("jiaxie/DeepProtT5-Fold", do_lower_case=False)
model = T5ClassificationModel.from_pretrained("jiaxie/DeepProtT5-Fold", torch_dtype=torch.bfloat16).to("cuda")