save snapshot

app.py

@@ -3,32 +3,47 @@ import gradio as gr
 from mammal_demo.demo_framework import MammalObjectBroker, MammalTask
 from mammal_demo.dti_task import DtiTask
 from mammal_demo.ppi_task import PpiTask
+from mammal_demo.tcr_task import TcrTask
 
 all_tasks: dict[str, MammalTask] = dict()
 all_models: dict[str, MammalObjectBroker] = dict()
 
+
+# first create the required tasks
+# Note that the tasks need access to the models, as the model to use depends on the state of the widget
+# we pass the all_models dict and update it when we actualy have the models.
 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
 
+tcr_task = TcrTask(model_dict=all_models)
+all_tasks[tcr_task.name] = tcr_task
+
+
+# create the model holders. hold the model and the tokenizer, lazy download
+# note that the list of relevent tasks needs to be stated.
 ppi_model = MammalObjectBroker(
-    model_path="ibm/biomed.omics.bl.sm.ma-ted-458m", task_list=[ppi_task.name]
+    model_path="ibm/biomed.omics.bl.sm.ma-ted-458m", task_list=[ppi_task.name,tcr_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],
+task_list=[tdi_task.name],
 )
 all_models[tdi_model.name] = tdi_model
 
+tcr_model = MammalObjectBroker(
+    model_path= "ibm/biomed.omics.bl.sm.ma-ted-458m.tcr_epitope_bind",
+    task_list=[tcr_task.name]
+)
+all_models[tcr_model.name] = tcr_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()

mammal_demo/ppi_task.py

@@ -140,9 +140,9 @@ class PpiTask(MammalTask):
                 )
             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")
+                score_box = gr.Number(label="PPI score")
                 run_mammal.click(
                     fn=self.create_and_run_prompt,
                     inputs=[model_name_widget, prot1, prot2],