Update tasks/audio.py

tasks/audio.py

@@ -55,36 +55,36 @@ async def evaluate_audio(request: AudioEvaluationRequest):
     
     feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base")
     # Prétraitement en streaming avec fonction explicite
-    test_dataset = test_dataset.map(preprocess_function, fn_kwargs={"feature_extractor": feature_extractor}, remove_columns="audio", batched=True)
+    test_dataset = test_dataset.map(preprocess_function, fn_kwargs={"feature_extractor": feature_extractor}, remove_columns="audio", batched=True,  batch_size=32)  # Choix de la taille du batch)
     
     gc.collect()
     
     # Pipeline de classification optimisé
     classifier = pipeline("audio-classification", model="CindyDelage/Challenge_HuggingFace_DFG_FrugalAI", device=device)
-    
-    logging.info("Début des prédictions par batch")
+     predictions = []  # Liste pour stocker les prédictions
 
-    with open('predictions.csv', mode='w', newline='') as file:
-        writer = csv.writer(file)
-        writer.writerow(['predicted_label'])  # Écrire les en-têtes
-        # Traiter les données et écrire les résultats dans le fichier
-        for data in test_dataset:
-            with torch.no_grad():
-                result = classifier(np.asarray(data["input_values"]), batch_size=2)
+    logging.info("Début des prédictions par batch")
+    for data in iter(test_dataset):
+        with torch.no_grad():
+            result = classifier(np.asarray(data["input_values"]), batch_size=2)
             
-            predicted_label = result[0]['label']
-            label = 1 if predicted_label == 'environment' else 0
-    
-            # Écrire chaque prédiction directement dans le fichier
-            writer.writerow([label])
-    
-            # Nettoyer la mémoire après chaque itération
-            del result
-            torch.cuda.empty_cache()
-            gc.collect()
+        predicted_label = result[0]['label']
+        label = 1 if predicted_label == 'environment' else 0
+        predictions.append(label)  # Ajouter la prédiction à la liste
+        
+        # Nettoyer la mémoire après chaque itération
+        del result
+        del label 
+        torch.cuda.empty_cache()
+        gc.collect()
         
     logging.info("Fin des prédictions")
+    del result
+    del label 
+    del classifier 
+    del feature_extractor 
     
+    gc.collect()
     # Stop tracking emissions
     emissions_data = tracker.stop_task()