Update app.py

app.py

@@ -5,11 +5,7 @@ from bs4 import BeautifulSoup
 import random
 
 classification_model = pipeline("text-classification", model="plantbert_text_classification_model", tokenizer="plantbert_text_classification_model")
-mask_model = pipeline("fill-mask", model="plantbert_fill_mask_model", tokenizer="plantbert_fill_mask_model", top_k=14189)
-
-def return_text(habitat_label):
-    text = f"This vegetation plot belongs to the habitat {habitat_label}."
-    return text
+mask_model = pipeline("fill-mask", model="plantbert_fill_mask_model", tokenizer="plantbert_fill_mask_model", top_k=100)
     
 def return_habitat_image(habitat_label):
     floraveg_url = f"https://floraveg.eu/habitat/overview/{habitat_label}"
@@ -23,7 +19,7 @@ def return_habitat_image(habitat_label):
             image_url = "https://www.salonlfc.com/wp-content/uploads/2018/01/image-not-found-scaled-1150x647.png"
     else:
         image_url = "https://www.salonlfc.com/wp-content/uploads/2018/01/image-not-found-scaled-1150x647.png"
-    image_url = "https://www.commissionoceanindien.org/wp-content/uploads/2018/07/plantnet.jpg"
+    image_url = "https://www.commissionoceanindien.org/wp-content/uploads/2018/07/plantnet.jpg"  # While we don't have the rights
     image = gr.Image(value=image_url)
     return image
 
@@ -70,9 +66,9 @@ def classification(text):
     text = gbif_normalization(text)
     result = classification_model(text)
     habitat_label = result[0]['label']
-    formatted_output = return_text(habitat_label)
+    text = f"This vegetation plot belongs to the habitat {habitat_label}."
     image_output = return_habitat_image(habitat_label)
-    return formatted_output, image_output
+    return text, image_output
 
 def masking(text):
     text = gbif_normalization(text)
@@ -141,7 +137,7 @@ def masking(text):
         best_position = len(text.split(', '))
         best_sentence = sentence
     
-    text = f"The most likely missing species is {best_prediction} at position {best_position}.\nThe new vegetation plot is {best_sentence}."
+    text = f"The most likely missing species is {best_prediction} (position {best_position}).\nThe new vegetation plot is {best_sentence}."
     image = return_species_image(best_prediction)
     return text, image