app.py

@@ -79,37 +79,50 @@ def classification(text, k):
     image_output = return_habitat_image(habitat_labels[0])
     return text, image_output
 
-def masking(text):
+def masking(text, k):
     text = gbif_normalization(text)
     text_split = text.split(', ')
     
-    max_score = 0
-    best_prediction = None
-    best_position = None
-    best_sentence = None
+    best_predictions = []
+    best_positions = []
+    best_sentences = []
+    
+    for _ in range(k):
+        max_score = 0
+        best_prediction = None
+        best_position = None
+        best_sentence = None
 
-    for i in range(len(text_split) + 1):
-        masked_text = ', '.join(text_split[:i] + ['[MASK]'] + text_split[i:])
-        
-        j = 0
-        while True:
-            prediction = mask_model(masked_text)[j]
-            species = prediction['token_str']
-            if species in text_split:
-                j += 1
-            else:
-                break
+        for i in range(len(text_split) + 1):
+            masked_text = ', '.join(text_split[:i] + ['[MASK]'] + text_split[i:])
+            
+            j = 0
+            while True:
+                prediction = mask_model(masked_text)[j]
+                species = prediction['token_str']
+                if species in text_split or species in best_predictions:
+                    j += 1
+                else:
+                    break
 
-        score = prediction['score']
-        sentence = prediction['sequence']
+            score = prediction['score']
+            sentence = prediction['sequence']
 
-        if score > max_score:
-            max_score = score
-            best_prediction = species
-            best_position = i
-            best_sentence = sentence
-    
-    text = f"The most likely missing species is {best_prediction} (position {best_position}).\nThe new vegetation plot is {best_sentence}."
+            if score > max_score:
+                max_score = score
+                best_prediction = species
+                best_position = i
+                best_sentence = sentence
+        
+        best_predictions.append(best_prediction)
+        best_positions.append(best_position)
+        best_sentences.append(best_sentence)
+        text_split.insert(best_position, best_prediction)
+    if k == 1:
+        text = f"The most likely missing species is {best_predictions[0]} (position {best_positions[0]})."
+    else:
+        text = f"The most likely missing species are {', '.join(best_predictions[:-1])} and {best_predictions[-1]} (positions {', '.join(map(str, best_positions[:-1]))} and {best_positions[-1]})."
+    text += f"\nThe new vegetation plot is {best_sentences[-1]}. (see image of the most likely species below)."
     image = return_species_image(best_prediction)
     return text, image
 
@@ -122,26 +135,27 @@ with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
                 species_classification = gr.Textbox(lines=2, label="Species", placeholder="Enter a list of comma-separated binomial names here.")
-                k_classification = gr.Slider(1, 5, value=1, label="Top-k", info="Choose the number of top habitats to display.")
+                k_classification = gr.Slider(1, 5, value=1, step=1, label="Top-k", info="Choose the number of top habitats to display.")
             with gr.Column():
-                text_output_1 = gr.Textbox()
-                text_output_2 = gr.Image()
+                text_classification = gr.Textbox()
+                image_classification = gr.Image()
         button_classification = gr.Button("Classify")
         gr.Markdown("""<h5 style="text-align: center;">An example of input</h5>""")
-        gr.Examples([["sparganium erectum, calystegia sepium, persicaria amphibia", 1]], [species_classification, k_classification], [text_output_1, text_output_2], classification, True)
+        gr.Examples([["sparganium erectum, calystegia sepium, persicaria amphibia", 1]], [species_classification, k_classification], [text_classification, image_classification], classification, True)
         
     with gr.Tab("Missing species finding"):
         gr.Markdown("""<h3 style="text-align: center;">Finding the missing species!</h3>""")
         with gr.Row():
             species_masking = gr.Textbox(lines=2, label="Species", placeholder="Enter a list of comma-separated binomial names here.")
+            k_masking = gr.Slider(1, 5, value=1, step=1, label="Top-k", info="Choose the number of top missing species to find.")
             with gr.Column():
-                image_output_1 = gr.Textbox()
-                image_output_2 = gr.Image()
+                text_masking = gr.Textbox()
+                image_masking = gr.Image()
         button_masking = gr.Button("Find")
         gr.Markdown("""<h5 style="text-align: center;">An example of input</h5>""")
-        gr.Examples([["vaccinium myrtillus, dryopteris dilatata, molinia caerulea"]], [species_masking], [image_output_1, image_output_2], masking, True)
+        gr.Examples([["vaccinium myrtillus, dryopteris dilatata, molinia caerulea", 1]], [species_masking, k_masking], [text_masking, image_masking], masking, True)
 
-    button_classification.click(classification, inputs=[species_classification, k_classification], outputs=[text_output_1, text_output_2])
-    button_masking.click(masking, inputs=[species_masking], outputs=[image_output_1, image_output_2])
+    button_classification.click(classification, inputs=[species_classification, k_classification], outputs=[textclassification, image_classification])
+    button_masking.click(masking, inputs=[species_masking, k_masking], outputs=[text_masking, image_masking])
 
 demo.launch()