tcr task works

mammal_demo/tcr_task.py

@@ -53,7 +53,7 @@ Given the TCR beta sequance and the epitope sequacne, estimate the binding speci
            dict: sample_dict for feeding into model
         """
         sample_dict= dict()
-        sample_dict[ENCODER_INPUTS_STR] = self.create_prompt(*sample_inputs)
+        sample_dict[ENCODER_INPUTS_STR] = self.create_prompt(**sample_inputs)
         tokenizer_op = model_holder.tokenizer_op
         model = model_holder.model
         tokenizer_op(
@@ -132,10 +132,11 @@ Given the TCR beta sequance and the epitope sequacne, estimate the binding speci
         else:
             scores=[None]
 
-        ans = dict(
-            pred=label_id_to_int.get(int(decoder_output[classification_position]), -1),
-            score=scores.item(),
-        )
+        ans = [
+            tokenizer_op._tokenizer.decode(batch_dict[CLS_PRED][0]),
+            label_id_to_int.get(int(decoder_output[classification_position]), -1),
+            scores.item(),
+        ]
         return ans
 
 
@@ -185,12 +186,13 @@ Given the TCR beta sequance and the epitope sequacne, estimate the binding speci
                 prompt_box = gr.Textbox(label="Mammal prompt", lines=5)
 
             with gr.Row():
-                decoded = gr.Textbox(label="Mammal prediction")
+                decoded = gr.Textbox(label="Mammal output")
+                predicted_class = gr.Textbox(label="Mammal prediction")
                 binding_score = gr.Number(label="Binding score")
                 run_mammal.click(
                     fn=self.create_and_run_prompt,
                     inputs=[model_name_widget, tcr_textbox, epitope_textbox],
-                    outputs=[prompt_box, decoded, binding_score],
+                    outputs=[prompt_box, decoded, predicted_class,binding_score],
                 )
             demo.visible = False
             return demo