Update tasks/audio.py

tasks/audio.py

@@ -3,12 +3,10 @@ from datetime import datetime
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score
 import numpy as np
-import random
 import os
 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
@@ -25,6 +23,8 @@ router = APIRouter()
 DESCRIPTION = "Random Baseline"
 ROUTE = "/audio"
 
+device = 0 if torch.cuda.is_available() else -1  # Choix du périphérique GPU si dispo
+
 @router.post(ROUTE, tags=["Audio Task"], description=DESCRIPTION)
 async def evaluate_audio(request: AudioEvaluationRequest):
     """
@@ -33,13 +33,14 @@ async def evaluate_audio(request: AudioEvaluationRequest):
     # Get space info
     username, space_url = get_space_info()
     
-    # Load dataset
+    # Load dataset en streaming
     logging.info("Chargement des données")
-    dataset = load_dataset(request.dataset_name, streaming=True,token=os.getenv("HF_TOKEN"))
+    dataset = load_dataset(request.dataset_name, streaming=True, token=os.getenv("HF_TOKEN"))
     logging.info("Données chargées")
     
     test_dataset = dataset["test"]
     del dataset
+    
     # Start tracking emissions
     tracker.start()
     tracker.start_task("inference")
@@ -47,39 +48,33 @@ async def evaluate_audio(request: AudioEvaluationRequest):
     # Feature extraction
     feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base")
     
-    def preprocess_function(examples):
-        audio_arrays = [x["array"] for x in examples["audio"]]
-        return feature_extractor(audio_arrays, sampling_rate=feature_extractor.sampling_rate, padding="longest", max_length=16000, truncation=True, return_tensors="pt")
+    def preprocess_function(example):
+        audio_array = example["audio"]["array"]
+        return feature_extractor(audio_array, sampling_rate=feature_extractor.sampling_rate, return_tensors="pt")
+    
+    # Prétraitement en streaming
+    test_dataset = test_dataset.map(preprocess_function, remove_columns=["audio"])
+    
+    del feature_extractor  # Libération explicite
+    gc.collect()
+    
+    # Pipeline de classification optimisé
+    classifier = pipeline("audio-classification", model="CindyDelage/Challenge_HuggingFace_DFG_FrugalAI", device=device)
     
-    encoded_data_test = test_dataset.map(preprocess_function, remove_columns="audio", batched=True)#, keep_in_memory=False)
-    del feature_extractor
-    del audio_arrays
-    # Pipeline de classification
-    classifier = pipeline("audio-classification", model="CindyDelage/Challenge_HuggingFace_DFG_FrugalAI", device=-1)
-
     predictions = []
     logging.info("Début des prédictions par batch")
-
-    for data in encoded_data_test:
-        # Récupérer les données audio et le label
+    
+    for data in test_dataset:
         with torch.no_grad():
             result = classifier(np.asarray(data["input_values"]))
         
         predicted_label = result[0]['label']
         predictions.append(1 if predicted_label == 'environment' else 0)
-        del result
-        del predicted_label
-
-        # Nettoyage mémoire après chaque batch
-        #del input_values
+        
+        del result  # Nettoyage mémoire
         torch.cuda.empty_cache()
         gc.collect()
         
-        # Log mémoire toutes les 500 prédictions
-        if len(predictions) % 500 == 0:
-            logging.info(f"Nettoyage mémoire après {len(predictions)} prédictions")
-            logging.info(f"Utilisation mémoire : {psutil.virtual_memory().percent}%")
-    
     logging.info("Fin des prédictions")
     
     # Stop tracking emissions
@@ -105,5 +100,6 @@ async def evaluate_audio(request: AudioEvaluationRequest):
             "test_seed": request.test_seed
         }
     }
+    
     logging.info("Returning results")
     return results