Update app.py

app.py

@@ -5,7 +5,7 @@ from bs4 import BeautifulSoup
 import random
 
 classification_model = pipeline("text-classification", model="plantbert_text_classification_model", tokenizer="plantbert_text_classification_model")
-mask_model = pipeline("fill-mask", model="plantbert_fill_mask_model", tokenizer="plantbert_fill_mask_model")
+mask_model = pipeline("fill-mask", model="plantbert_fill_mask_model", tokenizer="plantbert_fill_mask_model", top_k=14189)
 
 def return_text(habitat_label, habitat_score, confidence):
     if habitat_score*100 > confidence:
@@ -90,8 +90,14 @@ def masking(text):
 
     # Case for the first position
     masked_text = "[MASK], " + ', '.join(text.split(', '))
-    prediction = mask_model(masked_text)[0]
-    species = prediction['token_str']
+    i = 0
+    while True:
+        prediction = mask_model(masked_text)[i]
+        species = prediction['token_str']
+        if species in text.split(', '):
+            i+=1
+        else:
+            break
     score = prediction['score']
     sentence = prediction['sequence']
 
@@ -104,8 +110,14 @@ def masking(text):
     # Loop through each position in the middle of the sentence
     for i in range(1, len(text.split(', '))):
         masked_text = ', '.join(text.split(', ')[:i]) + ', [MASK], ' + ', '.join(text.split(', ')[i:])
-        prediction = mask_model(masked_text)[0]
-        species = prediction['token_str']
+        i = 0
+        while True:
+            prediction = mask_model(masked_text)[i]
+            species = prediction['token_str']
+            if species in text.split(', '):
+                i+=1
+            else:
+                break
         score = prediction['score']
         sentence = prediction['sequence']
         
@@ -118,8 +130,14 @@ def masking(text):
 
     # Case for the last position
     masked_text = ', '.join(text.split(', ')) + ', [MASK]'
-    prediction = mask_model(masked_text)[0]
-    species = prediction['token_str']
+    i = 0
+        while True:
+            prediction = mask_model(masked_text)[i]
+            species = prediction['token_str']
+            if species in text.split(', '):
+                i+=1
+            else:
+                break
     score = prediction['score']
     sentence = prediction['sequence']