Update app.py

app.py

@@ -57,8 +57,8 @@ def return_species_image(species):
     image = gr.Image(value=image_url)
     return image
 
-def classification(text, typology, confidence, task):
-    model = return_model(task)
+def classification(text, typology, confidence):
+    model = return_model("classification")
     dataset = return_dataset()
     result = model(text)
     habitat_label = result[0]['label']
@@ -68,8 +68,8 @@ def classification(text, typology, confidence, task):
     image_output = return_habitat_image(habitat_label, habitat_score, confidence)
     return formatted_output, image_output
 
-def masking(text, task):
-    model = return_model(task)
+def masking(text):
+    model = return_model("masking")
     masked_text = text + ', [MASK] [MASK]'
     pred = model(masked_text, top_k=1)
     new_species = [pred[i][0]['token_str'] for i in range(len(pred))]
@@ -84,14 +84,13 @@ with gr.Blocks() as demo:
         species = gr.Textbox(lines=2, label="Species", placeholder="Enter a list of comma-separated binomial names here.")
         typology = gr.Dropdown(["EUNIS"], value="EUNIS", label="Typology", info="Will add more typologies later!")
         confidence = gr.Slider(0, 100, value=90, label="Confidence", info="Choose the level of confidence for the prediction.")
-        task = gr.Radio(["classification", "masking"], value="classification", label="Task", info="Which task to choose?")
-        text_output_1 = gr.Textbox()
-        text_output_2 = gr.Image()
+        with gr.Row():
+            text_output_1 = gr.Textbox()
+            text_output_2 = gr.Image()
         text_button = gr.Button("Classify")
     with gr.Tab("Missing species finding"):
+        species_2 = gr.Textbox(lines=2, label="Species", placeholder="Enter a list of comma-separated binomial names here.")
         with gr.Row():
-            species_2 = gr.Textbox(lines=2, label="Species", placeholder="Enter a list of comma-separated binomial names here.")
-            task_2 = gr.Radio(["classification", "masking"], value="classification", label="Task", info="Which task to choose?")
             image_output_1 = gr.Textbox()
             image_output_2 = gr.Image()
         image_button = gr.Button("Find")