Update tasks/audio.py

tasks/audio.py

@@ -25,12 +25,12 @@ ROUTE = "/audio"
 
 device = 0 if torch.cuda.is_available() else -1  # Choix du périphérique GPU si dispo
 
- # 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, feature_extractor):
+    return feature_extractor(
+        example["audio"]["array"], 
+        sampling_rate=feature_extractor.sampling_rate, 
+        return_tensors="pt"
+    )
 
 @router.post(ROUTE, tags=["Audio Task"], description=DESCRIPTION)
 async def evaluate_audio(request: AudioEvaluationRequest):
@@ -56,7 +56,11 @@ async def evaluate_audio(request: AudioEvaluationRequest):
     feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base")
     
     # Prétraitement en streaming avec fonction explicite
-    encoded_data_test = test_dataset.map(preprocess_function, remove_columns="audio", batched=True)#, keep_in_memory=False)
+    test_dataset = test_dataset.map(
+        preprocess_function, 
+        fn_kwargs={"feature_extractor": feature_extractor}, 
+        remove_columns=["audio"]
+    )
     
     del feature_extractor  # Libération explicite
     gc.collect()
@@ -67,10 +71,9 @@ async def evaluate_audio(request: AudioEvaluationRequest):
     predictions = []
     logging.info("Début des prédictions par batch")
     
-    for data in encoded_data_test:
-        logging.info(data)
+    for data in test_dataset:
         with torch.no_grad():
-            result = classifier(np.asarray(data["input_values"]))
+            result = classifier(np.asarray(data["array"]))
         
         predicted_label = result[0]['label']
         predictions.append(1 if predicted_label == 'environment' else 0)