add audio task

tasks/audio.py

@@ -2,25 +2,34 @@ from fastapi import APIRouter
 from datetime import datetime
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score
+from accelerate import Accelerator
+from tqdm import tqdm
+from torch.amp import autocast
 import random
 import os
+import torch
 
 from .utils.evaluation import AudioEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
+from audio_utils import AudioClassifier, AudioDataset, Config, collate_fn, Evaluator
+from transformers import AutoModelForImageClassification
+from torch.utils.data import DataLoader
 
+from loguru import logger
 from dotenv import load_dotenv
 load_dotenv()
 
+
 router = APIRouter()
 
-DESCRIPTION = "Random Baseline"
+DESCRIPTION = "Audio pipeline to classify sounds."
 ROUTE = "/audio"
-
+device = "cuda"
 
 
 @router.post(ROUTE, tags=["Audio Task"],
              description=DESCRIPTION)
-async def evaluate_audio(request: AudioEvaluationRequest):
+async def evaluate_audio(request: AudioEvaluationRequest): #, model_path: str):
     """
     Evaluate audio classification for rainforest sound detection.
     
@@ -38,34 +47,44 @@ async def evaluate_audio(request: AudioEvaluationRequest):
     }
     # Load and prepare the dataset
     # Because the dataset is gated, we need to use the HF_TOKEN environment variable to authenticate
-    dataset = load_dataset(request.dataset_name,token=os.getenv("HF_TOKEN"))
+    config = Config()
+    accelerator = Accelerator()
+    device = accelerator.device
     
-    # Split dataset
-    train_test = dataset["train"]
+    dataset = load_dataset(request.dataset_name,token=os.getenv("HF_TOKEN"))
     test_dataset = dataset["test"]
-    
+    test_dataset = test_dataset.filter(lambda x: len(x["audio"]["array"]) > 0)
+    true_labels = test_dataset["label"]
+
+    test_dataset = AudioDataset(test_dataset)
+    test_loader = DataLoader(test_dataset, batch_size=2 * config.BATCH_SIZE, shuffle=False, collate_fn=collate_fn, num_workers=config.NUM_WORKERS, pin_memory=True)
+
+    model = AudioClassifier(config.MODEL_NAME, config.MODEL_PATH)
+    model, test_loader = accelerator.prepare(model, test_loader)
+
     # Start tracking emissions
     tracker.start()
     tracker.start_task("inference")
-    
+
     #--------------------------------------------------------------------------------------------
     # 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.
     #--------------------------------------------------------------------------------------------   
+    predictions = []
     
-    # Make random predictions (placeholder for actual model inference)
-    true_labels = test_dataset["label"]
-    predictions = [random.randint(0, 1) for _ in range(len(true_labels))]
+    logger.info("Running inference ...")
+    evaluator = Evaluator(model, test_loader, device)
+    predictions = evaluator.evaluate()
     
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE
     #--------------------------------------------------------------------------------------------   
-    
     # Stop tracking emissions
     emissions_data = tracker.stop_task()
     
     # Calculate accuracy
     accuracy = accuracy_score(true_labels, predictions)
+    print("accuracy", accuracy)
     
     # Prepare results dictionary
     results = {
@@ -84,5 +103,4 @@ async def evaluate_audio(request: AudioEvaluationRequest):
             "test_seed": request.test_seed
         }
     }
-    
-    return results 
+    return results