Update app.py

app.py

@@ -4,15 +4,25 @@ from datasets import load_dataset
 import requests
 from bs4 import BeautifulSoup
 
-classifier = pipeline("text-classification", model="CesarLeblanc/test_model")
-dataset = load_dataset("CesarLeblanc/text_classification_dataset")
+def return_model(task):
+    if task == 'classification':
+        model = pipeline("text-classification", model="CesarLeblanc/test_model")
+    else:
+        model = pipeline("fill-mask", model="CesarLeblanc/fill_mask_model")
+    return return_model
 
-def text_classification(text, typology, confidence):
-    result = classifier(text)
-    habitat_label = result[0]['label']
-    habitat_label = dataset['train'].features['label'].names[int(habitat_label.split('_')[1])]
-    habitat_score = result[0]['score']
-    formatted_output = f"This vegetation plot belongs to the habitat {habitat_label} with the probability {habitat_score*100:.2f}%"
+def return_dataset():
+    dataset = load_dataset("CesarLeblanc/text_classification_dataset")
+    return dataset
+
+def return_text(habitat_label, habitat_score, confidence):
+    if habitat_score*100 > confidence:
+        text = f"This vegetation plot belongs to the habitat {habitat_label} with the probability {habitat_score*100:.2f}%."
+    else:
+        text = f"We can't assign an habitat to this vegetation plot with a confidence of at least {confidence}%."
+    return text
+    
+def return_image(habitat_label, habitat_score, confidence):
     floraveg_url = f"https://floraveg.eu/habitat/overview/{habitat_label}"
     response = requests.get(floraveg_url)
     if response.status_code == 200:
@@ -21,22 +31,66 @@ def text_classification(text, typology, confidence):
         if img_tag:
             image_url = img_tag['src']
         else:
-            image_url = 'https://www.salonlfc.com/wp-content/uploads/2018/01/image-not-found-scaled-1150x647.png'
+            image_url = "https://www.salonlfc.com/wp-content/uploads/2018/01/image-not-found-scaled-1150x647.png"
+    else:
+        image_url = "https://www.salonlfc.com/wp-content/uploads/2018/01/image-not-found-scaled-1150x647.png"
+    if habitat_score*100 < confidence:
+        image_url = "https://www.salonlfc.com/wp-content/uploads/2018/01/image-not-found-scaled-1150x647.png"
+    image = gr.Image(value=image_url)
+    return image
+
+def classification(text, typology, confidence, task):
+    model = return_model(task)
+    dataset = return_dataset()
+    result = model(text)
+    habitat_label = result[0]['label']
+    habitat_label = dataset['train'].features['label'].names[int(habitat_label.split('_')[1])]
+    habitat_score = result[0]['score']
+    formatted_output = return_text(habitat_label, habitat_score, confidence)
+    image_output = return_image(habitat_label, habitat_score, confidence)
+    return formatted_output, image_output
+
+def masking(text, task):
+    model = return_model(task)
+    text += ', [MASK] [MASK]'
+    pred = mask_filler(text, top_k=1)
+    text = pred[0]["sequence"]
+    image = gr.Image(value="https://www.salonlfc.com/wp-content/uploads/2018/01/image-not-found-scaled-1150x647.png")
+    return text, image
+
+def plantbert(text, typology, confidence, task):
+    if task == "classification":
+        formatted_output, image_output = classification(text, typology, confidence, task)
     else:
-        image_url = 'https://www.salonlfc.com/wp-content/uploads/2018/01/image-not-found-scaled-1150x647.png'
-    image_output = gr.Image(value=image_url)
+        formatted_output, image_output = masking(text, typology, confidence, task)
     return formatted_output, image_output
 
+inputs=[
+    gr.Textbox(lines=2, label="Species", placeholder="Enter a list of comma-separated binomial names here."),
+    gr.Dropdown(["EUNIS"], value="EUNIS", label="Typology", info="Will add more typologies later!"),
+    gr.Slider(0, 100, value=90, label="Confidence", info="Choose the level of confidence for the prediction.")
+    gr.Radio(["classification", "masking"], value="classification", label="Task", info="Which task to choose?"),
+]
+
+outputs=[
+    gr.Textbox(lines=2, label="Vegetation Plot Classification Result"),
+    "image"
+]
+
+title="Pl@ntBERT"
+
+description="Vegetation Plot Classification: enter the species found in a vegetation plot and see its EUNIS habitat!"
+
 examples=[
-    ["sparganium erectum, calystegia sepium, persicaria amphibia", "EUNIS", 50],
-    ["thinopyrum junceum, cakile maritima", "EUNIS", 50]
+    ["sparganium erectum, calystegia sepium, persicaria amphibia", "EUNIS", 90, "classification"],
+    ["thinopyrum junceum, cakile maritima", "EUNIS", 90, "masking"]
 ]
 
-io = gr.Interface(fn=text_classification, 
-                         inputs=[gr.Textbox(lines=2, label="List of comma-separated binomial names of species (see examples)", placeholder="Enter species here..."), gr.Dropdown(["EUNIS"], label="Typology", info="Will add more typologies later!"), gr.Slider(0, 100, value=50, label="Confidence", info="Choose the level of confidence for the prediction")], 
-                         outputs=[gr.Textbox(lines=2, label="Vegetation Plot Classification Result"), "image"],
-                         title="Pl@ntBERT",
-                         description="Vegetation Plot Classification: enter the species found in a vegetation plot and see its EUNIS habitat!",
+io = gr.Interface(fn=plantbert, 
+                         inputs=inputs, 
+                         outputs=outputs,
+                         title=title,
+                         description=description,
                          examples=examples)
 
 io.launch()