added three MOLNET tasks. text in demo still not done

mammal_demo/__init__.py

@@ -7,7 +7,7 @@ from mammal_demo.dti_task import DtiTask
 from mammal_demo.ppi_task import PpiTask
 from mammal_demo.ps_task import PsTask
 from mammal_demo.tcr_task import TcrTask
-
+from mammal_demo.molnet_task import MolnetTask
 
 def tasks_and_models():
     all_tasks = TaskRegistry()
@@ -21,6 +21,9 @@ def tasks_and_models():
     tdi_task = all_tasks.register_task(DtiTask(model_dict=all_models))
     ps_task = all_tasks.register_task(PsTask(model_dict=all_models))
     tcr_task = all_tasks.register_task(TcrTask(model_dict=all_models))
+    bbbp_task = all_tasks.register_task(MolnetTask(model_dict=all_models,task_name="BBBP"))
+    toxicity_task = all_tasks.register_task(MolnetTask(model_dict=all_models,task_name="TOXICITY"))
+    fda_appr_task = all_tasks.register_task(MolnetTask(model_dict=all_models,task_name="FDA_APPR"))
 
 # create the model holders. hold the model and the tokenizer, lazy download
 # note that the list of relevent tasks needs to be stated.
@@ -46,13 +49,16 @@ def tasks_and_models():
     task_list=[ppi_task],
 )
     all_models.register_model(
-    "ibm/biomed.omics.bl.sm.ma-ted-458m.moleculenet_clintox_tox"
+    "ibm/biomed.omics.bl.sm.ma-ted-458m.moleculenet_clintox_tox",
+    task_list=[toxicity_task]
 )
     all_models.register_model(
-    "ibm/biomed.omics.bl.sm.ma-ted-458m.moleculenet_clintox_fda"
+    "ibm/biomed.omics.bl.sm.ma-ted-458m.moleculenet_clintox_fda",
+    task_list=[fda_appr_task]
 )
     all_models.register_model(
-    "ibm/biomed.omics.bl.sm.ma-ted-458m.moleculenet_bbbp"
+    "ibm/biomed.omics.bl.sm.ma-ted-458m.moleculenet_bbbp",
+    task_list=[bbbp_task],
 )
     
     return all_tasks,all_models

mammal_demo/molnet_task.py

@@ -0,0 +1,88 @@
+import gradio as gr
+from mammal.examples.molnet.molnet_infer import create_sample_dict as molnet_create_sample_dict, get_predictions, process_model_output
+from mammal.keys import *
+from mammal.model import Mammal
+
+from mammal_demo.demo_framework import MammalObjectBroker, MammalTask
+
+
+class MolnetTask(MammalTask):
+    def __init__(self, model_dict, task_name="BBBP"):
+        super().__init__(name=f"Molnet: {task_name}", model_dict=model_dict)
+        self.description = f"MOLNET {task_name}"
+        self.examples = {
+            "drug_seq": "CC(=O)NCCC1=CNc2c1cc(OC)cc2",
+        }
+        self.task_name=task_name
+        self.markup_text = """
+# Mammal based Drug-Target 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) -> dict:
+        return molnet_create_sample_dict(task_name=self.task_name, smiles_seq=sample_inputs["drug_seq"], tokenizer_op=model_holder.tokenizer_op, model=model_holder.model)
+
+    def run_model(self, sample_dict, model: Mammal):
+        # Generate Prediction
+        batch_dict = get_predictions(model=model,sample_dict=sample_dict)
+        return batch_dict
+
+    def decode_output(self, batch_dict, model_holder):
+        result = process_model_output(
+            tokenizer_op=model_holder.tokenizer_op,
+            decoder_output=batch_dict[CLS_PRED][0],
+            decoder_output_scores=batch_dict[SCORES][0],
+        )
+        generated_output = model_holder.tokenizer_op._tokenizer.decode(batch_dict[CLS_PRED][0])
+        return generated_output, result['pred'], result['score']
+
+    def create_and_run_prompt(self, model_name, drug_seq):
+        model_holder = self.model_dict[model_name]
+        inputs = {
+            "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():
+                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 task",
+                    variant="primary",
+                )
+            with gr.Row():
+                prompt_box = gr.Textbox(label="Mammal prompt", lines=5)
+
+            with gr.Row():
+                decoded = gr.Textbox(label="Mammal output")
+                prediction_box=gr.Textbox(label="Mammal prediction")
+                score_box=gr.Number(label="score")
+                run_mammal.click(
+                    fn=self.create_and_run_prompt,
+                    inputs=[model_name_widget, drug_textbox],
+                    outputs=[prompt_box, decoded, prediction_box, score_box],
+                )
+            demo.visible = False
+            return demo