update

app.py

@@ -18,7 +18,10 @@ model = AutoModelForCausalLM.from_pretrained(model_tokenizer_path, trust_remote_
 
 # List of available tasks
 tasks = ['H3', 'H4', 'H3K9ac', 'H3K14ac', 'H4ac', 'H3K4me1', 'H3K4me2', 'H3K4me3', 'H3K36me3', 'H3K79me3']
-
+mapping = {'1':'It is a',
+           '0':'It is not a',
+           'No valid prediction':'Cannot be determined whether or not it is a',
+          }
 def preprocess_response(response, mask_token="[MASK]"):
     """Extracts the response after the [MASK] token."""
     if mask_token in response:
@@ -48,8 +51,8 @@ def generate(dna_sequence, task_type, sample_num=1):
 
     response = model.generate(**tokenized_message, max_new_tokens=sample_num, do_sample=False)
     reply = tokenizer.batch_decode(response, skip_special_tokens=False)[0].replace(" ", "")
-
-    return extract_label(reply, task_type)
+    pred = extract_label(reply, task_type)
+    return f"{mapping[pred]} {task_type}"
 
 def extract_label(message, task_type):
     """Extracts the prediction label from the model's response."""
@@ -65,7 +68,7 @@ interface = gr.Interface(
         gr.Textbox(label="Input DNA Sequence", placeholder="Enter a DNA sequence"),
         gr.Dropdown(choices=tasks, label="Select Task Type"),
     ],
-    outputs=gr.Textbox(label="Predicted Function"),
+    outputs=gr.Textbox(label="Predicted Type"),
     title="Omni-DNA Multitask Prediction",
     description="Select a DNA-related task and input a sequence to generate function predictions.",
 )