Update app.py

app.py

@@ -86,17 +86,20 @@ def masking(text):
     max_score = 0
     best_prediction = None
     best_position = None
+    best_sentence = None
 
     # Case for the first position
     masked_text = "[MASK], " + ', '.join(text.split(', '))
     prediction = mask_model(masked_text)[0]
     species = prediction['token_str']
     score = prediction['score']
+    sentence = prediction['sequence']
 
     if score > max_score:
         max_score = score
         best_prediction = species
         best_position = 0
+        best_sentence = sentence
 
     # Loop through each position in the middle of the sentence
     for i in range(1, len(text.split(', '))):
@@ -104,25 +107,29 @@ def masking(text):
         prediction = mask_model(masked_text)[0]
         species = prediction['token_str']
         score = prediction['score']
+        sentence = prediction['sequence']
         
         # Update best prediction and position if score is higher
         if score > max_score:
             max_score = score
             best_prediction = species
             best_position = i
+            best_sentence = sentence
 
     # Case for the last position
     masked_text = ', '.join(text.split(', ')) + ', [MASK]'
     prediction = mask_model(masked_text)[0]
     species = prediction['token_str']
     score = prediction['score']
+    sentence = prediction['sequence']
     
     if score > max_score:
         max_score = score
         best_prediction = species
         best_position = len(text.split(', '))
+        best_sentence = sentence
     
-    text = f"The most likely missing species in position {best_position} is: {best_prediction}."
+    text = f"The most likely missing species is {best_prediction} at position {best_position}.\nThe new vegetation plot is {best_sentence}."
     image = return_species_image(best_prediction)
     return text, image