Update tasks/text.py

tasks/text.py

@@ -5,6 +5,9 @@ from sklearn.metrics import accuracy_score
 from sentence_transformers import SentenceTransformer
 import numpy as np
 
+from .utils.emissions import clean_emissions_data, get_space_info, tracker
+from .utils.evaluation import TextEvaluationRequest
+
 router = APIRouter()
 
 DESCRIPTION = "Class embeddings with cosine similarity using batching and thresholding"
@@ -42,10 +45,13 @@ class_embeddings = embedding_model.encode(class_labels, batch_size=8, convert_to
 
 
 @router.post(ROUTE, tags=["Text Task"], description=DESCRIPTION)
-async def evaluate_text(request: dict):
+async def evaluate_text(request: TextEvaluationRequest):
     """
     Evaluate text classification using precomputed embeddings and cosine similarity.
     """
+    # Get space info
+    username, space_url = get_space_info()
+
     # Load dataset
     dataset = load_dataset("QuotaClimat/frugalaichallenge-text-train")
     df_train = dataset["train"].to_pandas()
@@ -56,6 +62,14 @@ async def evaluate_text(request: dict):
     quotes = df["quote"].tolist()
     true_labels = df["label"].apply(lambda x: int(x.split("_")[0]) if isinstance(x, str) else 0).tolist()
 
+    # Start tracking emissions
+    tracker.start()
+    tracker.start_task("inference")
+
+    # --------------------------------------------------------------------------------------------
+    # Optimized cosine similarity-based classification
+    # --------------------------------------------------------------------------------------------
+
     # Encode dataset quotes in batches
     batch_size = 32
     quote_embeddings = embedding_model.encode(quotes, batch_size=batch_size, convert_to_numpy=True, normalize_embeddings=True)
@@ -70,14 +84,22 @@ async def evaluate_text(request: dict):
     # Apply threshold (0.9) for classification
     predicted_labels = [best_idx if best_sim > 0.9 else 0 for best_idx, best_sim in zip(best_indices, best_similarities)]
 
+    # Stop tracking emissions
+    emissions_data = tracker.stop_task()
+
     # Calculate accuracy
     accuracy = accuracy_score(true_labels, predicted_labels)
 
     # Prepare results dictionary
     results = {
+        "username": username,
+        "space_url": space_url,
         "submission_timestamp": datetime.now().isoformat(),
         "model_description": DESCRIPTION,
         "accuracy": float(accuracy),
+        "energy_consumed_wh": emissions_data.energy_consumed * 1000,
+        "emissions_gco2eq": emissions_data.emissions * 1000,
+        "emissions_data": clean_emissions_data(emissions_data),
         "api_route": ROUTE,
         "dataset_config": {
             "dataset_name": "QuotaClimat/frugalaichallenge-text-train",