matanninio's picture
refactoring to make code more elegant and cleanups
fda141d
raw
history blame
4.52 kB
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 sequance, estimate if it's soluble or insoluble.
"""
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 sequance",
# info="standard",
interactive=True,
lines=3,
value=self.examples["protein_seq"],
)
with gr.Row():
run_mammal = gr.Button(
"Run Mammal prompt for TCL-Epitope Interaction",
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 normelized score")
norm_score = gr.Number(label="Normelized 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