Update tasks/audio.py

tasks/audio.py

@@ -13,6 +13,7 @@ from .utils.emissions import tracker, clean_emissions_data, get_space_info
 from dotenv import load_dotenv
 import logging
 import csv
+import torch.nn.utils.prune as prune
 
 # Configurer le logging
 logging.basicConfig(level=logging.INFO)
@@ -32,6 +33,14 @@ def preprocess_function(example, feature_extractor):
         sampling_rate=feature_extractor.sampling_rate, padding="longest", max_length=16000, truncation=True, return_tensors="pt"
     )
 
+def apply_pruning(model, amount=0.3):
+    """Applique un pruning sur les poids du modèle."""
+    for name, module in model.named_modules():
+        if isinstance(module, torch.nn.Linear):
+            prune.l1_unstructured(module, name="weight", amount=amount)
+            prune.remove(module, "weight")
+    return model
+
 @router.post(ROUTE, tags=["Audio Task"], description=DESCRIPTION)
 async def evaluate_audio(request: AudioEvaluationRequest):
     """
@@ -51,28 +60,27 @@ async def evaluate_audio(request: AudioEvaluationRequest):
     tracker.start()
     tracker.start_task("inference")
 
-    
     feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base")
 
-    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)
+    test_dataset = test_dataset.map(preprocess_function, fn_kwargs={"feature_extractor": feature_extractor}, remove_columns="audio", batched=True, batch_size=32)  
     
     gc.collect()
 
-    # Charger le modèle depuis Hugging Face Hub (par exemple, à partir de l'ID du modèle)
-    model_name = "CindyDelage/Challenge_HuggingFace_DFG_FrugalAI"  # Nom du modèle dans Hugging Face Hub
-    model = Wav2Vec2ForSequenceClassification.from_pretrained(model_name)  # Charger le modèle pour la classification audio
-
-    # Appliquer la quantification dynamique si nécessaire
-    model.eval()  # Mettre le modèle en mode évaluation
-    model = torch.quantization.quantize_dynamic(model, dtype=torch.qint8)  # Appliquer la quantification dynamique
+    model_name = "CindyDelage/Challenge_HuggingFace_DFG_FrugalAI"
+    model = Wav2Vec2ForSequenceClassification.from_pretrained(model_name)
     
-    classifier = pipeline("audio-classification", model="CindyDelage/Challenge_HuggingFace_DFG_FrugalAI",feature_extractor=feature_extractor, device=device)
+    # Appliquer la quantification dynamique et le pruning
+    model.eval()
+    model = torch.quantization.quantize_dynamic(model, dtype=torch.qint8)
+    model = apply_pruning(model, amount=0.3)  # Prune 30% des poids linéaires
+    
+    classifier = pipeline("audio-classification", model=model, feature_extractor=feature_extractor, device=device)
     predictions = []  
     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=1)
-            
+        
         predicted_label = result[0]['label']
         label = 1 if predicted_label == 'environment' else 0
         predictions.append(label)  
@@ -84,8 +92,6 @@ async def evaluate_audio(request: AudioEvaluationRequest):
         gc.collect()
         
     logging.info("Fin des prédictions")
-    del result
-    del label 
     del classifier 
     del feature_extractor 
     
@@ -93,9 +99,7 @@ async def evaluate_audio(request: AudioEvaluationRequest):
     # Stop tracking emissions
     emissions_data = tracker.stop_task()
     
-    # Calculate accuracy
-    true_labels = []  # Liste pour stocker les labels réels
-
+    true_labels = []  
     for example in test_dataset:   
         true_labels.append(example["label"])   
 
@@ -119,4 +123,4 @@ async def evaluate_audio(request: AudioEvaluationRequest):
     }
     
     logging.info("Returning results")
-    return results
+    return results