Update app.py

app.py

@@ -4,16 +4,11 @@ from datasets import load_dataset
 import requests
 from bs4 import BeautifulSoup
 
-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 model
 
-def return_dataset():
-    dataset = load_dataset("CesarLeblanc/text_classification_dataset")
-    return dataset
+classification_model = pipeline("text-classification", model="CesarLeblanc/test_model")
+mask_model = pipeline("fill-mask", model="CesarLeblanc/fill_mask_model")
+
+dataset = load_dataset("CesarLeblanc/text_classification_dataset")
 
 def return_text(habitat_label, habitat_score, confidence):
     if habitat_score*100 > confidence:
@@ -58,9 +53,7 @@ def return_species_image(species):
     return image
 
 def classification(text, typology, confidence):
-    model = return_model("classification")
-    dataset = return_dataset()
-    result = model(text)
+    result = classification_model(text)
     habitat_label = result[0]['label']
     habitat_label = dataset['train'].features['label'].names[int(habitat_label.split('_')[1])]
     habitat_score = result[0]['score']
@@ -69,9 +62,8 @@ def classification(text, typology, confidence):
     return formatted_output, image_output
 
 def masking(text):
-    model = return_model("masking")
     masked_text = text + ', [MASK] [MASK]'
-    pred = model(masked_text, top_k=1)
+    pred = fill_model(masked_text, top_k=1)
     new_species = [pred[i][0]['token_str'] for i in range(len(pred))]
     new_species = ' '.join(new_species)
     text = f"The last species from this vegetation plot is probably {new_species}."