steveyu323/kinbert_v2_long

1from transformers import PreTrainedModel, PretrainedConfig2import torch3import torch.nn as nn4from tape import ProteinBertForSequenceClassification, TAPETokenizer5 6 7class KinaseSubstrateConfig(PretrainedConfig):8    model_type = "kinase_substrate_bert"9 10    def __init__(11        self,12        tape_model_name="bert-base",13        num_labels=2,14        threshold=0.5,15        with_sep=True,16        max_len=1024,17        **kwargs,18    ):19        super().__init__(**kwargs)20        self.tape_model_name = tape_model_name21        self.num_labels = num_labels22        self.threshold = threshold23        self.with_sep = with_sep24        self.max_len = max_len25 26 27class KinaseSubstrateModel(PreTrainedModel):28    config_class = KinaseSubstrateConfig29 30    def __init__(self, config: KinaseSubstrateConfig):31        super().__init__(config)32        self.backbone = ProteinBertForSequenceClassification.from_pretrained(33            config.tape_model_name, num_labels=config.num_labels34        )35 36    def forward(self, input_ids, input_mask=None, targets=None):37        return self.backbone(38            input_ids=input_ids,39            input_mask=input_mask,40            targets=targets,41        )42 43    def predict_proba(self, input_ids, input_mask):44        self.eval()45        with torch.no_grad():46            (_, _), logits = self.forward(input_ids=input_ids, input_mask=input_mask)47            return torch.softmax(logits, dim=-1)[:, 1]