Update tasks/text.py

tasks/text.py

@@ -2,6 +2,20 @@ from fastapi import APIRouter
 from datetime import datetime
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score
+
+from skops import hub_utils
+from skops.io import load
+
+from huggingface_hub import hf_hub_download
+import joblib
+
+REPO_ID = "kantundpeterpan/frugal-ai-toy"
+FILENAME = "tfidf_rf.pkl"
+
+model = joblib.load(
+    hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
+)
+
 import random
 
 from .utils.evaluation import TextEvaluationRequest
@@ -9,7 +23,7 @@ from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 router = APIRouter()
 
-DESCRIPTION = "Random Baseline"
+DESCRIPTION = "tfidf-rf"
 ROUTE = "/text"
 
 @router.post(ROUTE, tags=["Text Task"], 
@@ -55,10 +69,12 @@ async def evaluate_text(request: TextEvaluationRequest):
     # YOUR MODEL INFERENCE CODE HERE
     # Update the code below to replace the random baseline by your model inference within the inference pass where the energy consumption and emissions are tracked.
     #--------------------------------------------------------------------------------------------   
-    
+
     # Make random predictions (placeholder for actual model inference)
     true_labels = test_dataset["label"]
-    predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
+    predictions = [
+        LABEL_MAPPING[r] for r in model.predict(test_dataset)    
+    ]
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE