app.py

@@ -2,11 +2,10 @@ import gradio as gr
 from transformers import pipeline
 import requests
 from bs4 import BeautifulSoup
-import random
 
 # Initialize models
-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", top_k=100)
+classification_model = pipeline("text-classification", model="models/text_classification_model", tokenizer="models/text_classification_model")
+mask_model = pipeline("fill-mask", model="models/fill_mask_model", tokenizer="models/fill_mask_model", top_k=100)
     
 def return_habitat_image(habitat_label):
     floraveg_url = f"https://floraveg.eu/habitat/overview/{habitat_label}"
@@ -143,6 +142,7 @@ def masking(text):
     return text, image
 
 with gr.Blocks() as demo:
+
     gr.Markdown("""<h1 style="text-align: center;">Pl@ntBERT</h1>""")
     
     with gr.Tab("Vegetation plot classification"):
@@ -150,13 +150,13 @@ with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
                 species = gr.Textbox(lines=2, label="Species", placeholder="Enter a list of comma-separated binomial names here.")
-                top_k = 
+                top_k = gr.Slider(1, 5, value=1, label="Top-k", info="Choose between 1 and 5.")
             with gr.Column():
                 text_output_1 = gr.Textbox()
                 text_output_2 = gr.Image()
         text_button = gr.Button("Classify")
         gr.Markdown("""<h5 style="text-align: center;">An example of input</h5>""")
-        gr.Examples([["sparganium erectum, calystegia sepium, persicaria amphibia"]], [species], [text_output_1, text_output_2], classification, True)
+        gr.Examples([["sparganium erectum, calystegia sepium, persicaria amphibia", 1]], [species, top_k], [text_output_1, text_output_2], classification, True)
         
     with gr.Tab("Missing species finding"):
         gr.Markdown("""<h3 style="text-align: center;">Finding the missing species!</h3>""")
@@ -172,4 +172,4 @@ with gr.Blocks() as demo:
     text_button.click(classification, inputs=[species], outputs=[text_output_1, text_output_2])
     image_button.click(masking, inputs=[species_2], outputs=[image_output_1, image_output_2])
 
-demo.launch()
+demo.launch()