Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from fuse.data.tokenizers.modular_tokenizer.op import ModularTokenizerOp | |
| from mammal.examples.protein_solubility.task import ProteinSolubilityTask | |
| from mammal.keys import ( | |
| CLS_PRED, | |
| ENCODER_INPUTS_STR, | |
| SCORES, | |
| ) | |
| from mammal.model import Mammal | |
| from mammal_demo.demo_framework import MammalObjectBroker, MammalTask | |
| class PsTask(MammalTask): | |
| def __init__(self, model_dict): | |
| super().__init__(name="Protein Solubility", model_dict=model_dict) | |
| self.description = "Protein Solubility (PS)" | |
| self.examples = { | |
| "protein_seq": "LLQTGIHVRVSQPSL", | |
| } | |
| self.markup_text = """ | |
| # Mammal based protein solubility estimation | |
| Given the protein sequence, estimate if it's water-soluble. | |
| """ | |
| 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) # shallow copy | |
| sample_dict = ProteinSolubilityTask.data_preprocessing( | |
| sample_dict=sample_dict, | |
| protein_sequence_key="protein_seq", | |
| tokenizer_op=model_holder.tokenizer_op, | |
| device=model_holder.model.device, | |
| ) | |
| return sample_dict | |
| def run_model(self, sample_dict, model: Mammal): | |
| # Generate Prediction | |
| batch_dict = model.generate( | |
| [sample_dict], | |
| output_scores=True, | |
| return_dict_in_generate=True, | |
| max_new_tokens=5, | |
| ) | |
| return batch_dict | |
| def decode_output(self, batch_dict, tokenizer_op: ModularTokenizerOp) -> list: | |
| """ | |
| Extract predicted class and scores | |
| """ | |
| ans_dict = ProteinSolubilityTask.process_model_output( | |
| tokenizer_op=tokenizer_op, | |
| decoder_output=batch_dict[CLS_PRED][0], | |
| decoder_output_scores=batch_dict[SCORES][0], | |
| ) | |
| ans = [ | |
| tokenizer_op._tokenizer.decode(batch_dict[CLS_PRED][0]), | |
| ans_dict["pred"], | |
| ans_dict["not_normalized_scores"].item(), | |
| ans_dict["normalized_scores"].item(), | |
| ] | |
| return ans | |
| def create_and_run_prompt(self, model_name, protein_seq): | |
| model_holder = self.model_dict[model_name] | |
| inputs = { | |
| "protein_seq": protein_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, tokenizer_op=model_holder.tokenizer_op | |
| ) | |
| return res | |
| def create_demo(self, model_name_widget): | |
| with gr.Group() as demo: | |
| gr.Markdown(self.markup_text) | |
| with gr.Row(): | |
| protein_textbox = gr.Textbox( | |
| label="Protein sequence", | |
| # info="standard", | |
| interactive=True, | |
| lines=3, | |
| value=self.examples["protein_seq"], | |
| ) | |
| with gr.Row(): | |
| run_mammal = gr.Button( | |
| "Run Mammal prompt for protein solubility", | |
| variant="primary", | |
| ) | |
| with gr.Row(): | |
| prompt_box = gr.Textbox(label="Mammal prompt", lines=5) | |
| with gr.Row(): | |
| decoded = gr.Textbox(label="Mammal output") | |
| predicted_class = gr.Textbox(label="Mammal prediction") | |
| with gr.Column(): | |
| non_norm_score = gr.Number(label="Non normalized score") | |
| norm_score = gr.Number(label="normalized score") | |
| run_mammal.click( | |
| fn=self.create_and_run_prompt, | |
| inputs=[model_name_widget, protein_textbox], | |
| outputs=[ | |
| prompt_box, | |
| decoded, | |
| predicted_class, | |
| non_norm_score, | |
| norm_score, | |
| ], | |
| ) | |
| demo.visible = False | |
| return demo | |