Update tasks/text.py

tasks/text.py

@@ -13,7 +13,7 @@ from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 router = APIRouter()
 
-MODEL_TYPE = "modernbert_frugalai"
+MODEL_TYPE = "distilbert_frugalai"
 DESCRIPTIONS = {
     "distilbert_frugalai": "distilbert tuned on frugal ai data",
     "modernbert_frugalai": "distilbert tuned on frugal ai data"
@@ -137,7 +137,7 @@ async def evaluate_text(
     true_labels = test_dataset["label"]
     if model_type == "baseline":
         predictions = baseline_model(len(true_labels))
-    elif model_type[:5] == "bert-":
+    elif model_type in ["distilbert_frugalai", "modernbert_frugalai"]:
         predictions = bert_model(test_dataset, model_type)
     else:
         raise ValueError(model_type)