new_app.py

import gradio as gr
import torch
from fuse.data.tokenizers.modular_tokenizer.op import ModularTokenizerOp
from mammal.examples.dti_bindingdb_kd.task import DtiBindingdbKdTask
from mammal.keys import *
from mammal.model import Mammal
from abc import ABC, abstractmethod
class MammalObjectBroker():
    def __init__(self, model_path: str, name:str= None, task_list: list[str]=None) -> None:
        self.model_path = model_path
        if name is None:
            name = model_path
        self.name = name        
        
        if task_list is not None:
            self.tasks=task_list
        else:
            self.task = []
        self._model = None
        self._tokenizer_op = None
        
        
    @property
    def model(self)-> Mammal:
        if self._model is None:
            self._model =  Mammal.from_pretrained(self.model_path)
            self._model.eval()
        return self._model
    
    @property
    def tokenizer_op(self):
        if self._tokenizer_op is None:
            self._tokenizer_op =  ModularTokenizerOp.from_pretrained(self.model_path)
        return self._tokenizer_op
    
    
    
    

class MammalTask(ABC):
    def __init__(self, name:str) -> None:
            self.name = name
            self.description = None
            self._demo = None

    @abstractmethod
    def generate_prompt(self, **kwargs) -> str:
        """Formatting prompt to match pre-training syntax

        Args:
            prot1 (_type_): _description_
            prot2 (_type_): _description_

        Raises:
            No: _description_
        """
        raise NotImplementedError()

    @abstractmethod
    def crate_sample_dict(self, prompt: str, **kwargs) -> dict:
        """Formatting prompt to match pre-training syntax

        Args:
            prompt (str): _description_

        Returns:
            dict: sample_dict for feeding into model
        """
        raise NotImplementedError()

    # @abstractmethod
    def run_model(self, sample_dict, model:Mammal):
        raise NotImplementedError()
    
    @abstractmethod
    def create_demo(self, model_name_dropdown):
        """create an gradio demo group

        Returns:
            _type_: _description_
        """
        raise NotImplementedError()

    
    def demo(self,model_name_dropdown=None):
        if self._demo is None:
            self._demo = self.create_demo(model_name_dropdown=model_name_dropdown)
        return self._demo

    @abstractmethod
    def decode_output(self,batch_dict, model:Mammal):
        raise NotImplementedError()

    #self._setup()
        
    # def _setup(self):
    #     pass
    


all_tasks = dict()
all_models= dict()

class PpiTask(MammalTask):
    def __init__(self):
        super().__init__(name="PPI")
        self.description = "Protein-Protein Interaction (PPI)"
        self.examples = {
            "protein_calmodulin": "MADQLTEEQIAEFKEAFSLFDKDGDGTITTKELGTVMRSLGQNPTEAELQDMISELDQDGFIDKEDLHDGDGKISFEEFLNLVNKEMTADVDGDGQVNYEEFVTMMTSK",
            "protein_calcineurin": "MSSKLLLAGLDIERVLAEKNFYKEWDTWIIEAMNVGDEEVDRIKEFKEDEIFEEAKTLGTAEMQEYKKQKLEEAIEGAFDIFDKDGNGYISAAELRHVMTNLGEKLTDEEVDEMIRQMWDQNGDWDRIKELKFGEIKKLSAKDTRGTIFIKVFENLGTGVDSEYEDVSKYMLKHQ",
        }
        self.markup_text = """
    # Mammal based {self.description} demonstration
    
    Given two protein sequences, estimate if the proteins interact or not."""
    
        
        
    @staticmethod
    def positive_token_id(model_holder: MammalObjectBroker):
        """token for positive binding

        Args:
            model (MammalTrainedModel): model holding tokenizer

        Returns:
            int: id of positive binding token
        """
        return model_holder.tokenizer_op.get_token_id("<1>")
    
    def generate_prompt(self, prot1, prot2):
        """Formatting prompt to match pre-training syntax

        Args:
            prot1 (str): sequance of protein number 1
            prot2 (str): sequance of protein number 2

        Returns:
            str: prompt
        """   
        prompt =  "<@TOKENIZER-TYPE=AA><BINDING_AFFINITY_CLASS><SENTINEL_ID_0>"\
            "<MOLECULAR_ENTITY><MOLECULAR_ENTITY_GENERAL_PROTEIN>"\
            f"<SEQUENCE_NATURAL_START>{prot1}<SEQUENCE_NATURAL_END>"\
            "<MOLECULAR_ENTITY><MOLECULAR_ENTITY_GENERAL_PROTEIN>"\
            f"<SEQUENCE_NATURAL_START>{prot2}<SEQUENCE_NATURAL_END><EOS>"
        return prompt
    
    
    def crate_sample_dict(self,prompt: str, model_holder:MammalObjectBroker):
        # Create and load sample
        sample_dict = dict()
        sample_dict[ENCODER_INPUTS_STR] = prompt

        # Tokenize
        sample_dict = model_holder.tokenizer_op(
            sample_dict=sample_dict,
            key_in=ENCODER_INPUTS_STR,
            key_out_tokens_ids=ENCODER_INPUTS_TOKENS,
            key_out_attention_mask=ENCODER_INPUTS_ATTENTION_MASK,
        )
        sample_dict[ENCODER_INPUTS_TOKENS] = torch.tensor(
            sample_dict[ENCODER_INPUTS_TOKENS]
        )
        sample_dict[ENCODER_INPUTS_ATTENTION_MASK] = torch.tensor(
            sample_dict[ENCODER_INPUTS_ATTENTION_MASK]
        )
        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, model_holder):

        # Get output
        generated_output = model_holder.tokenizer_op._tokenizer.decode(batch_dict[CLS_PRED][0])
        score = batch_dict["model.out.scores"][0][1][self.positive_token_id(model_holder)].item()

        return generated_output, score


    def create_and_run_prompt(self,model_name,protein1, protein2):
        model_holder = all_models[model_name]
        prompt = self.generate_prompt(protein1, protein2)
        sample_dict = self.crate_sample_dict(prompt=prompt, model_holder=model_holder)
        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_dropdown):
        
    # """
    # ### Using the model from

    # ```{model} ```
    # """
        with gr.Group() as demo:
            gr.Markdown(self.markup_text)
            with gr.Row():
                prot1 = gr.Textbox(
                    label="Protein 1 sequence",
                    # info="standard",
                    interactive=True,
                    lines=3,
                    value=self.examples["protein_calmodulin"],
                )
                prot2 = gr.Textbox(
                    label="Protein 2 sequence",
                    # info="standard",
                    interactive=True,
                    lines=3,
                    value=self.examples["protein_calcineurin"],
                )
            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")
                run_mammal.click(
                    fn=self.create_and_run_prompt,
                    inputs=[model_name_dropdown, prot1, prot2],
                    outputs=[prompt_box, decoded, gr.Number(label="PPI score")],
                )
            with gr.Row():
                gr.Markdown(
                    "```<SENTINEL_ID_0>``` contains the binding affinity class, which is ```<1>``` for interacting and ```<0>``` for non-interacting"
                )
            demo.visible = True
            return demo

ppi_task = PpiTask()
all_tasks[ppi_task.name]=ppi_task

ppi_model = MammalObjectBroker(model_path="ibm/biomed.omics.bl.sm.ma-ted-458m", task_list=["PPI"])

all_models[ppi_model.name]=ppi_model
# tdi_model = MammalTrainedModel(model_path="ibm/biomed.omics.bl.sm.ma-ted-458m.dti_bindingdb_pkd")  TODO: ## task list still empty
# all_models.append(tdi_model)


def create_application():
    def task_change(value):
        choices=[model_name for model_name, model in all_models.items() if value in model.tasks]
        if choices:
            return  gr.update(choices=choices, value=choices[0])
        else:
            return
        # return model_name_dropdown
        
       
    with gr.Blocks() as demo:
        task_dropdown = gr.Dropdown(choices=["select demo"] + list(all_tasks.keys()))
        task_dropdown.interactive = True
        model_name_dropdown = gr.Dropdown(choices=[model_name for model_name, model in all_models.items() if task_dropdown.value in model.tasks], interactive=True)
        task_dropdown.change(task_change,inputs=[task_dropdown],outputs=[model_name_dropdown])
        
            



        ppi_demo = all_tasks["PPI"].demo(model_name_dropdown = model_name_dropdown)
        ppi_demo.visible = True
        # dtb_demo = create_tdb_demo()

        def set_ppi_vis(main_text):
            main_text=main_text
            print(f"main text is {main_text}")
            return gr.Group(visible=True)
            #return gr.Group(visible=(main_text == "PPI"))
        # , gr.Group(                visible=(main_text == "DTI")            )

        task_dropdown.change(
            set_ppi_vis, inputs=task_dropdown, outputs=[ppi_demo]
        )
        return demo

full_demo=None
def main():
    global full_demo
    full_demo = create_application()
    full_demo.launch(show_error=True, share=False)


if __name__ == "__main__":
    main()