Update tasks/audio.py

tasks/audio.py

@@ -5,18 +5,18 @@ from sklearn.metrics import accuracy_score
 import numpy as np
 import random
 import os
-from transformers import AutoFeatureExtractor
-
+import torch
+import gc
+import psutil
+from torch.utils.data import DataLoader
+from transformers import AutoFeatureExtractor, pipeline
 from .utils.evaluation import AudioEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
-
 from dotenv import load_dotenv
 import logging
 
 # Configurer le logging
 logging.basicConfig(level=logging.INFO)
-
-# Utiliser le logging au lieu de print
 logging.info("Début du fichier python")
 load_dotenv()
 
@@ -25,112 +25,69 @@ router = APIRouter()
 DESCRIPTION = "Random Baseline"
 ROUTE = "/audio"
 
-
-
-@router.post(ROUTE, tags=["Audio Task"],
-             description=DESCRIPTION)
-
+@router.post(ROUTE, tags=["Audio Task"], description=DESCRIPTION)
 async def evaluate_audio(request: AudioEvaluationRequest):
     """
     Evaluate audio classification for rainforest sound detection.
-    
-    Current Model: Random Baseline
-    - Makes random predictions from the label space (0-1)
-    - Used as a baseline for comparison
     """
     # Get space info
     username, space_url = get_space_info()
-
-    # Define the label mapping
-    LABEL_MAPPING = {
-        "chainsaw": 0,
-        "environment": 1
-    }
-    # Load and prepare the dataset
-    # Because the dataset is gated, we need to use the HF_TOKEN environment variable to authenticate
+    
+    # Load dataset
     logging.info("Chargement des données")
-    dataset = load_dataset(request.dataset_name,token=os.getenv("HF_TOKEN"))
+    dataset = load_dataset(request.dataset_name, token=os.getenv("HF_TOKEN"))
     logging.info("Données chargées")
-    # Split dataset
-    train_test = dataset["train"]
+    
     test_dataset = dataset["test"]
     
     # 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.
-    #--------------------------------------------------------------------------------------------   
-    
-    # Make random predictions (placeholder for actual model inference)
-    true_labels = test_dataset["label"]
-    import torch
-    from transformers import pipeline
-    from sklearn import preprocessing
-    from transformers import AutoFeatureExtractor
-
+    # Feature extraction
     feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base")
-
+    
     def preprocess_function(examples):
         audio_arrays = [x["array"] for x in examples["audio"]]
-        inputs = feature_extractor(audio_arrays, sampling_rate=feature_extractor.sampling_rate,  padding="longest", max_length=16000, truncation=True,return_tensors="pt")
-        return inputs
-   
+        return feature_extractor(audio_arrays, sampling_rate=feature_extractor.sampling_rate, padding="longest", max_length=16000, truncation=True, return_tensors="pt")
+    
     encoded_data_test = test_dataset.map(preprocess_function, remove_columns="audio", batched=True)
     
-    from datasets import Dataset
-    from transformers import AutoFeatureExtractor
-
-    # Utilisation du pipeline directement sur le dataset
-    classifier = pipeline("audio-classification", 
-                      model="CindyDelage/Challenge_HuggingFace_DFG_FrugalAI", 
-                      device=-1)
-    # Correctly access the audio data
-    # audio_data = [example["array"] for example in dataset["test"]["audio"]]
+    # Pipeline de classification
+    classifier = pipeline("audio-classification", model="CindyDelage/Challenge_HuggingFace_DFG_FrugalAI", device=-1)
+    
+    # DataLoader pour batch processing
+    BATCH_SIZE = 8
+    dataset_for_loader = [{"input_values": torch.tensor(example["input_values"])} for example in encoded_data_test]
+    dataloader = DataLoader(dataset_for_loader, batch_size=BATCH_SIZE)
+    
     predictions = []
-
-    logging.info("Début des prédictions")
+    logging.info("Début des prédictions par batch")
     
-    for example in encoded_data_test:
-        logging.info("Nombre de prédictions faites :", len(predictions))
-        input_values = np.array(example["input_values"])
-        result = classifier(input_values)  # Utilisation des données pré-traitées
-        #result = classifier(example["input_values"])  # Utilisation des données pré-traitées
-        predicted_label = result[0]['label']
-        predictions.append(1 if predicted_label == 'environment' else 0)
-
-
-    logging.info("Fin des prédictions")
-    #predictions = []
-    # for result in results:
-        # Check if result is a dictionary
-         #if isinstance(result, dict):
-    #        # Get the label with the highest score
-     #       predicted_label = result['label']
-      #  else:
-            # If result is not a dictionary, access it as a list
-       #     predicted_label = result[0]['label']  # Assuming the dictionary is the first element
+    for batch in dataloader:
+        input_values = batch["input_values"]
+        results = classifier(input_values)  # Pipeline en batch
+        
+        for result in results:
+            predicted_label = result[0]['label']
+            predictions.append(1 if predicted_label == 'environment' else 0)
+        
+        # Nettoyage mémoire toutes les 500 prédictions
+        if len(predictions) % 500 == 0:
+            torch.cuda.empty_cache()
+            gc.collect()
+            logging.info(f"Nettoyage de la mémoire après {len(predictions)} prédictions")
+            logging.info(f"Utilisation mémoire : {psutil.virtual_memory().percent}%")
     
-        # Assign 1 for "environment", 0 for "chainsaw"
-        #if predicted_label == 'environment':
-         #   predictions.append(1)
-        #else:
-         #   predictions.append(0)
-        #print(len(predictions))
-
-    #--------------------------------------------------------------------------------------------
-    # YOUR MODEL INFERENCE STOPS HERE
-    #--------------------------------------------------------------------------------------------   
+    logging.info("Fin des prédictions")
     
     # Stop tracking emissions
     emissions_data = tracker.stop_task()
     
     # Calculate accuracy
+    true_labels = test_dataset["label"]
     accuracy = accuracy_score(true_labels, predictions)
     
-    # Prepare results dictionary
     results = {
         "username": username,
         "space_url": space_url,
@@ -148,4 +105,4 @@ async def evaluate_audio(request: AudioEvaluationRequest):
         }
     }
     logging.info("Returning results")
-    return results 
+    return results