Update app.py

app.py

@@ -10,6 +10,8 @@ masking_model = pipeline("fill-mask", model="models/fill_mask_model", tokenizer=
 
 eunis_habitats = pd.read_excel('data/eunis_habitats.xlsx')
 
+image_not_found = gr.Image("https://img.freepik.com/premium-vector/file-folder-mascot-character-design-vector_166742-4413.jpg")
+
 def return_image(task, label):
     image = None
     if task == "classification":
@@ -61,7 +63,7 @@ def classification(text, k):
         text += "\nBelow is an example of this habitat type taken from the website FloraVEG."
     else:
         text += f"\nNo image found for this habitat type."
-        image = gr.Image("https://img.freepik.com/premium-vector/file-folder-mascot-character-design-vector_166742-4413.jpg")
+        image = image_not_found
     return text, image
 
 def masking(text, k):
@@ -97,15 +99,12 @@ def masking(text, k):
                 best_position = i
                 best_sentence = sentence
         
-        best_predictions.append(best_prediction.capitalize())
+        best_predictions.append(best_prediction)
         text_split.insert(best_position, best_prediction)
         
     best_positions = [text_split.index(prediction) for prediction in best_predictions]
-
-    
-    best_sentence = ", ".join(
-        [s.strip().capitalize() for s in best_sentence.split(",")]
-    )
+    best_predictions = [s.strip().capitalize() for s in best_predictions]
+    best_sentence = ", ".join([s.strip().capitalize() for s in best_sentence.split(",")])
     
     if k == 1:
         text = f"The most likely missing species is {best_predictions} (position {best_positions[0]})."
@@ -119,7 +118,7 @@ def masking(text, k):
         text += f"\nBelow is an image of the first missing species (i.e., {best_predictions}) taken from the website FloraVEG."
     else:
         text += f"\nNo image found for the first missing species (i.e., {best_predictions})."
-        image = gr.Image("https://img.freepik.com/premium-vector/file-folder-mascot-character-design-vector_166742-4413.jpg")
+        image = image_not_found
     return text, image
 
 with gr.Blocks() as demo: