Update README.md

README.md

@@ -1,3 +1,58 @@
----
-license: apache-2.0
----
+---
+license: apache-2.0
+---
+
+First we define a class T5RegressionModel:
+```python
+from transformers import (
+    T5Config,
+    T5EncoderModel,
+    T5Tokenizer,
+    PreTrainedModel,
+    TrainingArguments,
+    Trainer,
+    DataCollatorWithPadding,
+)
+class T5RegressionModel(PreTrainedModel):
+
+    config_class = T5Config
+
+    def __init__(self, config, d_model=None):
+        super().__init__(config)
+
+        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.regression_head = nn.Linear(hidden_dim, 1)
+
+    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.regression_head(pooled_output).squeeze(-1)  # [batch_size]
+
+
+        loss = None
+        if labels is not None:
+            labels = labels.float()
+            loss = nn.MSELoss()(logits, labels)
+
+        return {
+            "loss": loss,
+            "logits": logits
+        }
+```
+Then we load our pretrained model
+```python
+tokenizer = T5Tokenizer.from_pretrained("jiaxie/DeepProtT5-Stability", do_lower_case=False)
+model = T5RegressionModel.from_pretrained("jiaxie/DeepProtT5-Stability", torch_dtype=torch.bfloat16).to("cuda")
+```