Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from mammal.examples.dti_bindingdb_kd.task import DtiBindingdbKdTask | |
| from mammal.keys import * | |
| from mammal.model import Mammal | |
| from mammal_demo.demo_framework import MammalObjectBroker, MammalTask | |
| class DtiTask(MammalTask): | |
| def __init__(self, model_dict): | |
| super().__init__(name="Drug-Target Binding Affinity", model_dict=model_dict) | |
| self.description = "Drug-Target Binding Affinity (tdi)" | |
| self.examples = { | |
| "target_seq": "NLMKRCTRGFRKLGKCTTLEEEKCKTLYPRGQCTCSDSKMNTHSCDCKSC", | |
| "drug_seq": "CC(=O)NCCC1=CNc2c1cc(OC)cc2", | |
| } | |
| self.markup_text = """ | |
| # Mammal based Target-Drug binding affinity demonstration | |
| Given a protein sequence and a drug (in SMILES), estimate the binding affinity. | |
| """ | |
| def crate_sample_dict(self, sample_inputs: dict, model_holder: MammalObjectBroker): | |
| """convert sample_inputs to sample_dict including creating a proper prompt | |
| Args: | |
| sample_inputs (dict): dictionary containing the inputs to the model | |
| model_holder (MammalObjectBroker): model holder | |
| Returns: | |
| dict: sample_dict for feeding into model | |
| """ | |
| sample_dict = dict(sample_inputs) | |
| sample_dict = DtiBindingdbKdTask.data_preprocessing( | |
| sample_dict=sample_dict, | |
| tokenizer_op=model_holder.tokenizer_op, | |
| target_sequence_key="target_seq", | |
| drug_sequence_key="drug_seq", | |
| norm_y_mean=None, | |
| norm_y_std=None, | |
| device=model_holder.model.device, | |
| ) | |
| return sample_dict | |
| def run_model(self, sample_dict, model: Mammal): | |
| # Generate Prediction | |
| batch_dict = model.forward_encoder_only([sample_dict]) | |
| return batch_dict | |
| def decode_output(self, batch_dict, model_holder): | |
| # Get output | |
| batch_dict = DtiBindingdbKdTask.process_model_output( | |
| batch_dict, | |
| scalars_preds_processed_key="model.out.dti_bindingdb_kd", | |
| norm_y_mean=5.79384684128215, | |
| norm_y_std=1.33808027428196, | |
| ) | |
| ans = ( | |
| "model.out.dti_bindingdb_kd", | |
| float(batch_dict["model.out.dti_bindingdb_kd"][0]), | |
| ) | |
| return ans | |
| def create_and_run_prompt(self, model_name, target_seq, drug_seq): | |
| model_holder = self.model_dict[model_name] | |
| inputs = { | |
| "target_seq": target_seq, | |
| "drug_seq": drug_seq, | |
| } | |
| sample_dict = self.crate_sample_dict( | |
| sample_inputs=inputs, model_holder=model_holder | |
| ) | |
| prompt = sample_dict[ENCODER_INPUTS_STR] | |
| batch_dict = self.run_model(sample_dict=sample_dict, model=model_holder.model) | |
| res = prompt, *self.decode_output(batch_dict, model_holder=model_holder) | |
| return res | |
| def create_demo(self, model_name_widget): | |
| # """ | |
| # ### Using the model from | |
| # ```{model} ``` | |
| # """ | |
| with gr.Group() as demo: | |
| gr.Markdown(self.markup_text) | |
| with gr.Row(): | |
| target_textbox = gr.Textbox( | |
| label="target sequence", | |
| # info="standard", | |
| interactive=True, | |
| lines=3, | |
| value=self.examples["target_seq"], | |
| ) | |
| drug_textbox = gr.Textbox( | |
| label="Drug sequance (in SMILES)", | |
| # info="standard", | |
| interactive=True, | |
| lines=3, | |
| value=self.examples["drug_seq"], | |
| ) | |
| with gr.Row(): | |
| run_mammal = gr.Button( | |
| "Run Mammal prompt for Protein-Protein Interaction", | |
| variant="primary", | |
| ) | |
| with gr.Row(): | |
| prompt_box = gr.Textbox(label="Mammal prompt", lines=5) | |
| with gr.Row(): | |
| decoded = gr.Textbox(label="Mammal output key") | |
| run_mammal.click( | |
| fn=self.create_and_run_prompt, | |
| inputs=[model_name_widget, target_textbox, drug_textbox], | |
| outputs=[prompt_box, decoded, gr.Number(label="binding affinity")], | |
| ) | |
| demo.visible = False | |
| return demo | |