| from transformers import PreTrainedModel, PretrainedConfig |
| import torch |
| import torch.nn as nn |
| from tape import ProteinBertForSequenceClassification, TAPETokenizer |
|
|
|
|
| class KinaseSubstrateConfig(PretrainedConfig): |
| model_type = "kinase_substrate_bert" |
|
|
| def __init__( |
| self, |
| tape_model_name="bert-base", |
| num_labels=2, |
| threshold=0.5, |
| with_sep=True, |
| max_len=1024, |
| **kwargs, |
| ): |
| super().__init__(**kwargs) |
| self.tape_model_name = tape_model_name |
| self.num_labels = num_labels |
| self.threshold = threshold |
| self.with_sep = with_sep |
| self.max_len = max_len |
|
|
|
|
| class KinaseSubstrateModel(PreTrainedModel): |
| config_class = KinaseSubstrateConfig |
|
|
| def __init__(self, config: KinaseSubstrateConfig): |
| super().__init__(config) |
| self.backbone = ProteinBertForSequenceClassification.from_pretrained( |
| config.tape_model_name, num_labels=config.num_labels |
| ) |
|
|
| def forward(self, input_ids, input_mask=None, targets=None): |
| return self.backbone( |
| input_ids=input_ids, |
| input_mask=input_mask, |
| targets=targets, |
| ) |
|
|
| def predict_proba(self, input_ids, input_mask): |
| self.eval() |
| with torch.no_grad(): |
| (_, _), logits = self.forward(input_ids=input_ids, input_mask=input_mask) |
| return torch.softmax(logits, dim=-1)[:, 1] |
