app.py

@@ -4,7 +4,7 @@ import requests
 from bs4 import BeautifulSoup
 
 # Initialize models
-classification_model = pipeline("text-classification", model="models/text_classification_model", tokenizer="models/text_classification_model")
+classification_model = pipeline("text-classification", model="models/text_classification_model", tokenizer="models/text_classification_model", top_k=5)
 mask_model = pipeline("fill-mask", model="models/fill_mask_model", tokenizer="models/fill_mask_model", top_k=100)
     
 def return_habitat_image(habitat_label):
@@ -65,9 +65,12 @@ def gbif_normalization(text):
 def classification(text, k):
     text = gbif_normalization(text)
     result = classification_model(text)
-    habitat_label = result[0]['label']
-    text = f"This vegetation plot belongs to the habitat {habitat_label}."
-    image_output = return_habitat_image(habitat_label)
+    habitat_labels = [res['label'] for res in result[:k]]
+    if k == 1:
+        text = f"This vegetation plot belongs to the habitat {habitat_labels[0]}."
+    else:
+        text = f"This vegetation plot belongs to the habitats {', '.join(habitat_labels[:-1])} and {habitat_labels[-1]}."
+    image_output = return_habitat_image(habitat_labels[0])
     return text, image_output
 
 def masking(text):