File size: 1,432 Bytes
e1bca72 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | 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] |