Views
No views yet
1from transformers import (
2 T5Config,
3 T5EncoderModel,
4 T5Tokenizer,
5 PreTrainedModel,
6 TrainingArguments,
7 Trainer,
8 DataCollatorWithPadding,
9)
10class T5ClassificationModel(PreTrainedModel):
11 config_class = T5Config
12
13 def __init__(self, config, d_model=None, num_classes=1195):
14 super().__init__(config)
15 self.num_classes = num_classes
16
17 self.encoder = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_uniref50")
18
19 hidden_dim = d_model if d_model is not None else config.d_model
20 self.classification_head = nn.Linear(hidden_dim, num_classes)
21
22 def forward(
23 self,
24 input_ids=None,
25 attention_mask=None,
26 labels=None,
27 **kwargs
28 ):
29 encoder_outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
30 hidden_states = encoder_outputs.last_hidden_state
31
32 mask = attention_mask.unsqueeze(-1)
33 pooled_output = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1)
34 logits = self.classification_head(pooled_output) # [batch_size, num_classes]
35
36 loss = None
37 if labels is not None:
38 labels = labels.to(torch.bfloat16)
39 loss = nn.CrossEntropyLoss()(logits, labels)
40
41 return {
42 "loss": loss,
43 "logits": logits
44 }1tokenizer = T5Tokenizer.from_pretrained("jiaxie/DeepProtT5-Fold", do_lower_case=False)
2model = T5ClassificationModel.from_pretrained("jiaxie/DeepProtT5-Fold", torch_dtype=torch.bfloat16).to("cuda")