| 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, |
| } |
