File size: 1,936 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 48 49 50 51 52 53 54 55 | from transformers import Pipeline
from tape import TAPETokenizer
import torch
import numpy as np
class KinaseSubstratePipeline(Pipeline):
"""
Usage:
pipe = pipeline(
"kinase-substrate",
model="your-username/kinase-substrate-classifier",
trust_remote_code=True,
)
result = pipe({"kinase_seq": "MGSSHHH...", "substrate_seq": "ARTKQTAR..."})
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.tape_tokenizer = TAPETokenizer(vocab="iupac")
def _sanitize_parameters(self, **kwargs):
return {}, {}, {}
def preprocess(self, inputs):
kinase_seq = inputs["kinase_seq"].upper().replace(" ", "")
substrate_seq = inputs["substrate_seq"].upper().replace(" ", "")
cfg = self.model.config
kin_toks = self.tape_tokenizer.tokenize(kinase_seq)
sub_toks = self.tape_tokenizer.tokenize(substrate_seq)
toks = kin_toks + (["<sep>"] if cfg.with_sep else []) + sub_toks
toks = self.tape_tokenizer.add_special_tokens(toks)
ids = self.tape_tokenizer.convert_tokens_to_ids(toks)[: cfg.max_len]
input_ids = torch.tensor([ids], dtype=torch.long)
input_mask = torch.ones_like(input_ids)
return {"input_ids": input_ids, "input_mask": input_mask}
def _forward(self, model_inputs):
prob = self.model.predict_proba(
input_ids=model_inputs["input_ids"].to(self.device),
input_mask=model_inputs["input_mask"].to(self.device),
)
return {"prob": prob.cpu().numpy()}
def postprocess(self, model_outputs):
prob = float(model_outputs["prob"][0])
threshold = self.model.config.threshold
return {
"probability": round(prob, 4),
"label": "interaction" if prob >= threshold else "no_interaction",
"threshold": threshold,
} |