Update app.py

app.py

@@ -52,10 +52,10 @@ def classification(text, typology, confidence, task):
 
 def masking(text, task):
     model = return_model(task)
-    text += ', [MASK] [MASK]'
-    pred = model(text, top_k=1)
-    text = pred["sequence"]
-    #text = pred[0]["sequence"]
+    masked_text = text + ', [MASK] [MASK]'
+    pred = model(masked_text, top_k=1)
+    new_species = [pred[i][0]['token_str'] for i in range(len(pred))]
+    text = text + ', ' + ' '.join(new_species)
     image = gr.Image(value="https://www.salonlfc.com/wp-content/uploads/2018/01/image-not-found-scaled-1150x647.png")
     return text, image