Update app.py

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 from transformers import pipeline
-from datasets import load_dataset
 import requests
 from bs4 import BeautifulSoup
 import random
@@ -84,7 +83,8 @@ def classification(text, typology, confidence):
 def masking(text):
     text = gbif_normalization(text)
     masked_text = text + ', [MASK]'
-    new_species = mask_model(masked_text)[0]
+    pred = mask_model(masked_text)[0]
+    new_species = pred['token_str']
     text = f"The last species from this vegetation plot is probably {new_species}."
     image = return_species_image(new_species)
     return text, image