improve submission

tasks/text.py

@@ -1,5 +1,5 @@
 # tasks/text.py
-from fastapi import APIRouter
+from fastapi import APIRouter, HTTPException
 from datetime import datetime
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score
@@ -9,7 +9,7 @@ from torch.utils.data import Dataset, DataLoader
 import logging
 
 from .utils.evaluation import TextEvaluationRequest
-from .utils.emissions import tracker, clean_emissions_data, get_space_info
+from .utils.emissions import start_tracking, stop_tracking, clean_emissions_data, get_space_info
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -70,27 +70,22 @@ async def evaluate_text(request: TextEvaluationRequest):
         }
 
         logger.info("Loading dataset")
-        # Load dataset
         dataset = load_dataset(request.dataset_name)
-        
-        # Convert string labels to integers
         dataset = dataset.map(lambda x: {"label": LABEL_MAPPING[x["label"]]})
-        
-        # Get test dataset
         test_dataset = dataset["test"]
         
         logger.info("Starting emissions tracking")
-        # Start tracking emissions
-        tracker.start()
+        start_tracking()
         
         try:
-            # Load model and tokenizer
             logger.info("Loading model and tokenizer")
             model_name = "Tonic/climate-guard-toxic-agent"
             tokenizer = AutoTokenizer.from_pretrained(model_name)
-            model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=len(LABEL_MAPPING))
+            model = AutoModelForSequenceClassification.from_pretrained(
+                model_name, 
+                num_labels=len(LABEL_MAPPING)
+            )
             
-            # Prepare dataset
             logger.info("Preparing dataset")
             test_data = TextDataset(
                 texts=test_dataset["text"],
@@ -100,7 +95,6 @@ async def evaluate_text(request: TextEvaluationRequest):
             
             test_loader = DataLoader(test_data, batch_size=16)
             
-            # Model inference
             device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
             logger.info(f"Using device: {device}")
             model = model.to(device)
@@ -122,14 +116,11 @@ async def evaluate_text(request: TextEvaluationRequest):
                     predictions.extend(predicted.cpu().numpy())
                     ground_truth.extend(labels.cpu().numpy())
             
-            # Calculate accuracy
             accuracy = accuracy_score(ground_truth, predictions)
             logger.info(f"Accuracy: {accuracy}")
             
-            # Stop tracking emissions
-            emissions_data = tracker.stop()
+            emissions_data = stop_tracking()
             
-            # Prepare results
             results = {
                 "username": username,
                 "space_url": space_url,
@@ -138,7 +129,7 @@ async def evaluate_text(request: TextEvaluationRequest):
                 "accuracy": float(accuracy),
                 "energy_consumed_wh": float(emissions_data.energy_consumed * 1000),
                 "emissions_gco2eq": float(emissions_data.emissions * 1000),
-                "emissions_data": clean_emissions_data(emissions_data),
+                "emissions_data": clean_emissions_data(emissions_data.__dict__),
                 "api_route": ROUTE,
                 "dataset_config": {
                     "dataset_name": request.dataset_name,
@@ -152,9 +143,4 @@ async def evaluate_text(request: TextEvaluationRequest):
             
         except Exception as e:
             logger.error(f"Error during evaluation: {str(e)}")
-            tracker.stop()
-            raise e
-            
-    except Exception as e:
-        logger.error(f"Error in evaluate_text: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
+            stop_tracking()

tasks/utils/emissions.py

@@ -1,8 +1,13 @@
+# tasks/utils/emissions.py
 from codecarbon import EmissionsTracker
 import os
 
-# Initialize tracker
-tracker = EmissionsTracker(allow_multiple_runs=True)
+# Initialize tracker with correct parameters
+tracker = EmissionsTracker(
+    project_name="frugal-ai-challenge",
+    output_dir=".",
+    log_level='warning'
+)
 
 class EmissionsData:
     def __init__(self, energy_consumed: float, emissions: float):
@@ -11,8 +16,15 @@ class EmissionsData:
 
 def clean_emissions_data(emissions_data):
     """Remove unwanted fields from emissions data"""
-    data_dict = emissions_data.__dict__
-    fields_to_remove = ['timestamp', 'project_name', 'experiment_id', 'latitude', 'longitude']
+    if isinstance(emissions_data, dict):
+        data_dict = emissions_data
+    else:
+        data_dict = emissions_data.__dict__
+    
+    fields_to_remove = [
+        'timestamp', 'project_name', 'experiment_id', 
+        'latitude', 'longitude', '_start_time', '_end_time'
+    ]
     return {k: v for k, v in data_dict.items() if k not in fields_to_remove}
 
 def get_space_info():
@@ -25,4 +37,23 @@ def get_space_info():
             return username, space_url
         except Exception as e:
             print(f"Error getting space info: {e}")
-    return "local-user", "local-development" 
+    return "local-user", "local-development"
+
+def start_tracking():
+    """Start tracking emissions"""
+    try:
+        tracker.start()
+    except Exception as e:
+        print(f"Error starting emissions tracking: {e}")
+
+def stop_tracking():
+    """Stop tracking emissions and return data"""
+    try:
+        emissions = tracker.stop()
+        return EmissionsData(
+            energy_consumed=emissions.energy_consumed,
+            emissions=emissions.emissions
+        )
+    except Exception as e:
+        print(f"Error stopping emissions tracking: {e}")
+        return EmissionsData(energy_consumed=0.0, emissions=0.0)