Update app.py

app.py

@@ -87,8 +87,11 @@ def masking(text):
     text = gbif_normalization(text)
     masked_text = text + ', [MASK] [MASK]'
     pred = mask_model(masked_text, top_k=1)
-    new_species = [pred[i][0]['token_str'] for i in range(len(pred))]
-    new_species = ' '.join(new_species)
+    new_genus = pred[0][0]['token_str']
+    masked_text = text + f', {new_genus} [MASK]'
+    pred = mask_model(masked_text, top_k=1)
+    new_epithet = pred[0]['token_str']
+    new_species = new_genus + ' ' + new_epithet
     text = f"The last species from this vegetation plot is probably {new_species}."
     image = return_species_image(new_species)
     return text, image