cleanup and minor touches + renamed to the standard app name

app.py

@@ -1,247 +1,86 @@
 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
 
-model_paths = dict()
-
-ppi = "Protein-Protein Interaction (PPI)"
-model_paths[ppi] = "ibm/biomed.omics.bl.sm.ma-ted-458m"
-
-#
-dti = "Drug-Target Binding Affinity"
-model_paths[dti] = "ibm/biomed.omics.bl.sm.ma-ted-458m.dti_bindingdb_pkd"
-
-
-# load models (should probably be lazy)
-
-models = dict()
-tokenizer_op = dict()
-
-
-for task, model_path in model_paths.items():
-    if task not in models:
-        models[task] = Mammal.from_pretrained(model_path)
-        models[task].eval()
-        # Load Tokenizer
-        tokenizer_op[task] = ModularTokenizerOp.from_pretrained(model_path)
-
-
-### PPI:
-# token for positive binding
-positive_token_id = tokenizer_op[ppi].get_token_id("<1>")
-
-# Default input proteins
-protein_calmodulin = "MADQLTEEQIAEFKEAFSLFDKDGDGTITTKELGTVMRSLGQNPTEAELQDMISELDQDGFIDKEDLHDGDGKISFEEFLNLVNKEMTADVDGDGQVNYEEFVTMMTSK"
-protein_calcineurin = "MSSKLLLAGLDIERVLAEKNFYKEWDTWIIEAMNVGDEEVDRIKEFKEDEIFEEAKTLGTAEMQEYKKQKLEEAIEGAFDIFDKDGNGYISAAELRHVMTNLGEKLTDEEVDEMIRQMWDQNGDWDRIKELKFGEIKKLSAKDTRGTIFIKVFENLGTGVDSEYEDVSKYMLKHQ"
-
-
-def format_prompt_ppi(prot1, prot2):
-    # Formatting prompt to match pre-training syntax
-    return f"<@TOKENIZER-TYPE=AA><BINDING_AFFINITY_CLASS><SENTINEL_ID_0><MOLECULAR_ENTITY><MOLECULAR_ENTITY_GENERAL_PROTEIN><SEQUENCE_NATURAL_START>{prot1}<SEQUENCE_NATURAL_END><MOLECULAR_ENTITY><MOLECULAR_ENTITY_GENERAL_PROTEIN><SEQUENCE_NATURAL_START>{prot2}<SEQUENCE_NATURAL_END><EOS>"
-
-
-def run_prompt(prompt):
-    # Create and load sample
-    sample_dict = dict()
-    sample_dict[ENCODER_INPUTS_STR] = prompt
-
-    # Tokenize
-    sample_dict = tokenizer_op[ppi](
-        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]
-    )
-
-    # Generate Prediction
-    batch_dict = models[ppi].generate(
-        [sample_dict],
-        output_scores=True,
-        return_dict_in_generate=True,
-        max_new_tokens=5,
-    )
-
-    # Get output
-    generated_output = tokenizer_op[ppi]._tokenizer.decode(batch_dict[CLS_PRED][0])
-    score = batch_dict["model.out.scores"][0][1][positive_token_id].item()
-
-    return generated_output, score
-
-
-def create_and_run_prompt(protein1, protein2):
-    prompt = format_prompt_ppi(protein1, protein2)
-    res = prompt, *run_prompt(prompt=prompt)
-    return res
-
-
-def create_ppi_demo():
-    markup_text = f"""
-# Mammal based Protein-Protein Interaction (PPI) demonstration
-
-Given two protein sequences, estimate if the proteins interact or not.
-
-### Using the model from
-
- ```{model_paths[ppi]} ```
-"""
-    with gr.Group() as ppi_demo:
-        gr.Markdown(markup_text)
-        with gr.Row():
-            prot1 = gr.Textbox(
-                label="Protein 1 sequence",
-                # info="standard",
-                interactive=True,
-                lines=3,
-                value=protein_calmodulin,
-            )
-            prot2 = gr.Textbox(
-                label="Protein 2 sequence",
-                # info="standard",
-                interactive=True,
-                lines=3,
-                value=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=create_and_run_prompt,
-                inputs=[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"
-            )
-        ppi_demo.visible = False
-        return ppi_demo
-
-
-### DTI:
-# input
-target_seq = "NLMKRCTRGFRKLGKCTTLEEEKCKTLYPRGQCTCSDSKMNTHSCDCKSC"
-drug_seq = "CC(=O)NCCC1=CNc2c1cc(OC)cc2"
-
-
-# token for positive binding
-positive_token_id = tokenizer_op[dti].get_token_id("<1>")
-
-
-def format_prompt_dti(prot, drug):
-    sample_dict = {"target_seq": target_seq, "drug_seq": drug_seq}
-    sample_dict = DtiBindingdbKdTask.data_preprocessing(
-        sample_dict=sample_dict,
-        tokenizer_op=tokenizer_op[dti],
-        target_sequence_key="target_seq",
-        drug_sequence_key="drug_seq",
-        norm_y_mean=None,
-        norm_y_std=None,
-        device=models[dti].device,
-    )
-    return sample_dict
-
-
-def create_and_run_prompt_dtb(prot, drug):
-    sample_dict = format_prompt_dti(prot, drug)
-    # Post-process the model's output
-    # batch_dict = model_dti.forward_encoder_only([sample_dict])
-    batch_dict = models[dti].forward_encoder_only([sample_dict])
-    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]),
-    ]
-    res = sample_dict["data.query.encoder_input"], *ans
-    return res
-
-
-def create_tdb_demo():
-    markup_text = f"""
-# Mammal based Target-Drug binding affinity demonstration
-
-Given a protein sequence and a drug (in SMILES), estimate the binding affinity.
-
-### Using the model from
-
- ```{model_paths[dti]} ```
-"""
-    with gr.Group() as tdb_demo:
-        gr.Markdown(markup_text)
-        with gr.Row():
-            prot = gr.Textbox(
-                label="Protein sequence",
-                # info="standard",
-                interactive=True,
-                lines=3,
-                value=target_seq,
-            )
-            drug = gr.Textbox(
-                label="drug sequence (SMILES)",
-                # info="standard",
-                interactive=True,
-                lines=3,
-                value=drug_seq,
-            )
-        with gr.Row():
-            run_mammal = gr.Button(
-                "Run Mammal prompt for Target Drug Affinity", 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=create_and_run_prompt_dtb,
-                inputs=[prot, drug],
-                outputs=[prompt_box, decoded, gr.Number(label="DTI score")],
-            )
-        tdb_demo.visible = False
-        return tdb_demo
+from mammal_demo.demo_framework import MammalObjectBroker, MammalTask
+from mammal_demo.dti_task import DtiTask
+from mammal_demo.ppi_task import PpiTask
 
+all_tasks: dict[str, MammalTask] = dict()
+all_models: dict[str, MammalObjectBroker] = dict()
 
-def create_application():
+ppi_task = PpiTask(model_dict=all_models)
+all_tasks[ppi_task.name] = ppi_task
+
+tdi_task = DtiTask(model_dict=all_models)
+all_tasks[tdi_task.name] = tdi_task
+
+ppi_model = MammalObjectBroker(
+    model_path="ibm/biomed.omics.bl.sm.ma-ted-458m", task_list=[ppi_task.name]
+)
+all_models[ppi_model.name] = ppi_model
 
-    with gr.Blocks() as demo:
-        main_dropdown = gr.Dropdown(choices=["select demo", ppi, dti])
-        main_dropdown.interactive = True
-        ppi_demo = create_ppi_demo()
-        dtb_demo = create_tdb_demo()
+tdi_model = MammalObjectBroker(
+    model_path="ibm/biomed.omics.bl.sm.ma-ted-458m.dti_bindingdb_pkd",
+    task_list=[tdi_task.name],
+)
+all_models[tdi_model.name] = tdi_model
 
-        def set_ppi_vis(main_text):
-            return gr.Group(visible=main_text == ppi), gr.Group(
-                visible=main_text == dti
-            )
 
-        main_dropdown.change(
-            set_ppi_vis, inputs=main_dropdown, outputs=[ppi_demo, dtb_demo]
+def create_application():
+    def task_change(value):
+        visibility = [gr.update(visible=(task == value)) for task in all_tasks.keys()]
+        # all_tasks[task].demo().visible =
+        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], visible=True), *visibility)
+        else:
+            return (gr.skip, *visibility)
+        # return model_name_dropdown
+
+    with gr.Blocks() as application:
+        task_dropdown = gr.Dropdown(choices=["select demo"] + list(all_tasks.keys()), label="Mammal Task")
+        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,
+            label="Matching Mammal models",
+            visible=False,
+        )
+
+        task_dropdown.change(
+            task_change,
+            inputs=[task_dropdown],
+            outputs=[model_name_dropdown]
+            + [all_tasks[task].demo(model_name_widgit=model_name_dropdown) for task in all_tasks],
         )
-        return demo
+
+        # def set_demo_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 application
+
+
+full_demo = None
 
 
 def main():
-    demo = create_application()
-    demo.launch(show_error=True, share=True)
+    global full_demo
+    full_demo = create_application()
+    full_demo.launch(show_error=True, share=False)
 
 
 if __name__ == "__main__":

mammal_demo/demo_framework.py

@@ -1,51 +1,48 @@
+from abc import ABC, abstractmethod
+
 import gradio as gr
 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:
+class MammalObjectBroker:
+    def __init__(
+        self,
+        model_path: str,
+        name: str | None = None,
+        task_list: list[str] | None = None,
+    ) -> None:
         self.model_path = model_path
         if name is None:
             name = model_path
-        self.name = name        
-        
+        self.name = name
+
+        self.tasks: list[str] = []
         if task_list is not None:
-            self.tasks=task_list
-        else:
-            self.task = []
-        self._model = None
+            self.tasks = task_list
+        self._model: Mammal | None = None
         self._tokenizer_op = None
-        
-        
+
     @property
-    def model(self)-> Mammal:
+    def model(self) -> Mammal:
         if self._model is None:
-            self._model =  Mammal.from_pretrained(self.model_path)
-            self._model.eval()
+            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)
+            self._tokenizer_op = ModularTokenizerOp.from_pretrained(self.model_path)
         return self._tokenizer_op
-    
-    
-    
+
 
 class MammalTask(ABC):
-    def __init__(self, name:str, model_dict: dict[str,MammalObjectBroker]) -> None:
-            self.name = name
-            self.description = None
-            self._demo = None
-            self.model_dict = model_dict
+    def __init__(self, name: str, model_dict: dict[str, MammalObjectBroker]) -> None:
+        self.name = name
+        self.description = None
+        self._demo = None
+        self.model_dict = model_dict
 
     # @abstractmethod
     # def _generate_prompt(self, **kwargs) -> str:
@@ -61,7 +58,9 @@ class MammalTask(ABC):
     #     raise NotImplementedError()
 
     @abstractmethod
-    def crate_sample_dict(self,sample_inputs: dict, model_holder:MammalObjectBroker) -> dict:
+    def crate_sample_dict(
+        self, sample_inputs: dict, model_holder: MammalObjectBroker
+    ) -> dict:
         """Formatting prompt to match pre-training syntax
 
         Args:
@@ -73,9 +72,9 @@ class MammalTask(ABC):
         raise NotImplementedError()
 
     # @abstractmethod
-    def run_model(self, sample_dict, model:Mammal):
+    def run_model(self, sample_dict, model: Mammal):
         raise NotImplementedError()
-    
+
     def create_demo(self, model_name_widget: gr.component) -> gr.Group:
         """create an gradio demo group
 
@@ -89,20 +88,17 @@ class MammalTask(ABC):
         """
         raise NotImplementedError()
 
-
-    
-    def demo(self,model_name_widgit:gr.component=None):
+    def demo(self, model_name_widgit: gr.component = None):
         if self._demo is None:
-            model_name_widget:gr.component
+            model_name_widget: gr.component
             self._demo = self.create_demo(model_name_widget=model_name_widgit)
         return self._demo
 
     @abstractmethod
-    def decode_output(self,batch_dict, model:Mammal):
+    def decode_output(self, batch_dict, model: Mammal):
         raise NotImplementedError()
 
-    #self._setup()
-        
+    # self._setup()
+
     # def _setup(self):
     #     pass
-    

mammal_demo/dti_task.py

@@ -3,7 +3,8 @@ 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    
+from mammal_demo.demo_framework import MammalObjectBroker, MammalTask
+
 
 class DtiTask(MammalTask):
     def __init__(self, model_dict):
@@ -11,15 +12,15 @@ class DtiTask(MammalTask):
         self.description = "Drug-Target Binding Affinity (tdi)"
         self.examples = {
             "target_seq": "NLMKRCTRGFRKLGKCTTLEEEKCKTLYPRGQCTCSDSKMNTHSCDCKSC",
-            "drug_seq":"CC(=O)NCCC1=CNc2c1cc(OC)cc2"
-            }
+            "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):
+
+    def crate_sample_dict(self, sample_inputs: dict, model_holder: MammalObjectBroker):
         """convert sample_inputs to sample_dict including creating a proper prompt
 
         Args:
@@ -39,14 +40,13 @@ Given a protein sequence and a drug (in SMILES), estimate the binding affinity.
             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):
+
+    def decode_output(self, batch_dict, model_holder):
 
         # Get output
         batch_dict = DtiBindingdbKdTask.process_model_output(
@@ -54,34 +54,34 @@ Given a protein sequence and a drug (in SMILES), estimate the binding affinity.
             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]),
-        ) 
+            "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):
+    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]
+        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)
+        res = prompt, *self.decode_output(batch_dict, model_holder=model_holder)
         return res
 
-    
-    def create_demo(self,model_name_widget):
-        
-    # """
-    # ### Using the model from
+    def create_demo(self, model_name_widget):
+
+        # """
+        # ### Using the model from
 
-    # ```{model} ```
-    # """
+        # ```{model} ```
+        # """
         with gr.Group() as demo:
             gr.Markdown(self.markup_text)
             with gr.Row():
@@ -101,7 +101,8 @@ Given a protein sequence and a drug (in SMILES), estimate the binding affinity.
                 )
             with gr.Row():
                 run_mammal = gr.Button(
-                    "Run Mammal prompt for Protein-Protein Interaction", variant="primary"
+                    "Run Mammal prompt for Protein-Protein Interaction",
+                    variant="primary",
                 )
             with gr.Row():
                 prompt_box = gr.Textbox(label="Mammal prompt", lines=5)

mammal_demo/ppi_task.py

@@ -1,12 +1,14 @@
 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.keys import (
+    CLS_PRED,
+    ENCODER_INPUTS_ATTENTION_MASK,
+    ENCODER_INPUTS_STR,
+    ENCODER_INPUTS_TOKENS,
+)
 from mammal.model import Mammal
 
-from mammal_demo.demo_framework import MammalObjectBroker, MammalTask    
-
+from mammal_demo.demo_framework import MammalObjectBroker, MammalTask
 
 
 class PpiTask(MammalTask):
@@ -19,11 +21,9 @@ class PpiTask(MammalTask):
         }
         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
@@ -35,7 +35,7 @@ class PpiTask(MammalTask):
             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
 
@@ -45,16 +45,17 @@ class PpiTask(MammalTask):
 
         Returns:
             str: prompt
-        """   
-        prompt =  f"<@TOKENIZER-TYPE=AA><BINDING_AFFINITY_CLASS><SENTINEL_ID_0>"\
-            "<MOLECULAR_ENTITY><MOLECULAR_ENTITY_GENERAL_PROTEIN>"\
-            "<SEQUENCE_NATURAL_START>{prot1}<SEQUENCE_NATURAL_END>"\
-            "<MOLECULAR_ENTITY><MOLECULAR_ENTITY_GENERAL_PROTEIN>"\
-            "<SEQUENCE_NATURAL_START>{prot2}<SEQUENCE_NATURAL_END><EOS>"
+        """
+        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,sample_inputs: dict, model_holder:MammalObjectBroker):
+
+    def crate_sample_dict(self, sample_inputs: dict, model_holder: MammalObjectBroker):
         # Create and load sample
         sample_dict = dict()
         prompt = self.generate_prompt(*sample_inputs)
@@ -84,35 +85,37 @@ class PpiTask(MammalTask):
             max_new_tokens=5,
         )
         return batch_dict
-        
-    def decode_output(self,batch_dict, model_holder:MammalObjectBroker):
+
+    def decode_output(self, batch_dict, model_holder: MammalObjectBroker):
 
         # 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()
+        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):
+    def create_and_run_prompt(self, model_name, protein1, protein2):
         model_holder = self.model_dict[model_name]
-        sample_inputs = {"prot1":protein1, 
-                  "prot2":protein2
-                  }
-        sample_dict = self.crate_sample_dict(sample_inputs=sample_inputs, model_holder=model_holder)
+        sample_inputs = {"prot1": protein1, "prot2": protein2}
+        sample_dict = self.crate_sample_dict(
+            sample_inputs=sample_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)
+        res = prompt, *self.decode_output(batch_dict, model_holder=model_holder)
         return res
 
-    
-    def create_demo(self,model_name_widget:gr.component):
-        
-    # """
-    # ### Using the model from
+    def create_demo(self, model_name_widget: gr.component):
+
+        # """
+        # ### Using the model from
 
-    # ```{model} ```
-    # """
+        # ```{model} ```
+        # """
         with gr.Group() as demo:
             gr.Markdown(self.markup_text)
             with gr.Row():
@@ -132,17 +135,18 @@ class PpiTask(MammalTask):
                 )
             with gr.Row():
                 run_mammal: gr.Button = gr.Button(
-                    "Run Mammal prompt for Protein-Protein Interaction", variant="primary"
+                    "Run Mammal prompt for Protein-Protein Interaction",
+                    variant="primary",
                 )
             with gr.Row():
                 prompt_box = gr.Textbox(label="Mammal prompt", lines=5)
-
+                score_box = gr.Number(label="PPI score")
             with gr.Row():
                 decoded = gr.Textbox(label="Mammal output")
                 run_mammal.click(
                     fn=self.create_and_run_prompt,
                     inputs=[model_name_widget, prot1, prot2],
-                    outputs=[prompt_box, decoded, gr.Number(label="PPI score")],
+                    outputs=[prompt_box, decoded, score_box],
                 )
             with gr.Row():
                 gr.Markdown(

new_app.py

@@ -1,76 +0,0 @@
-import gradio as gr
-from mammal.keys import *
-
-from mammal_demo.demo_framework import MammalObjectBroker    
-
-
-from mammal_demo.ppi_task import PpiTask
-from mammal_demo.dti_task import DtiTask
-
-all_tasks = dict()
-all_models= dict()
-
-ppi_task = PpiTask(model_dict = all_models)
-all_tasks[ppi_task.name]=ppi_task
-
-tdi_task = DtiTask(model_dict = all_models)
-all_tasks[tdi_task.name]=tdi_task
-
-ppi_model = MammalObjectBroker(model_path="ibm/biomed.omics.bl.sm.ma-ted-458m", task_list=[ppi_task.name])
-all_models[ppi_model.name]=ppi_model
-
-tdi_model = MammalObjectBroker(model_path="ibm/biomed.omics.bl.sm.ma-ted-458m.dti_bindingdb_pkd", task_list=[tdi_task.name])
-all_models[tdi_model.name]=tdi_model
-
-
-def create_application():
-    def task_change(value):
-        visibility = [gr.update(visible=(task==value)) for task in all_tasks.keys()]
-            # all_tasks[task].demo().visible = 
-        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]),*visibility)
-        else:
-            return (gr.skip,*visibility)
-        # return model_name_dropdown
-        
-       
-    with gr.Blocks() as application:
-        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)
-        
-            
-
-
-
-        ppi_demo = all_tasks[ppi_task.name].demo(model_name_widgit = model_name_dropdown)
-        # ppi_demo.visible = True
-        dtb_demo = all_tasks[tdi_task.name].demo(model_name_widgit = model_name_dropdown)
-
-        task_dropdown.change(task_change,inputs=[task_dropdown],outputs=[model_name_dropdown]+[all_tasks[task].demo() for task in all_tasks])
-        
-        # def set_demo_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 application
-
-full_demo=None
-
-def main():
-    global full_demo
-    full_demo = create_application()
-    full_demo.launch(show_error=True, share=False)
-
-
-if __name__ == "__main__":
-    main()
-