Update tools.py

tools.py

@@ -51,66 +51,71 @@ class AudioEvaluationRequest(BaseEvaluationRequest):
     dataset_name: str = Field("rfcx/frugalai", 
     description="The name of the dataset on HuggingFace Hub") 
     
-def evaluate_audio(request: AudioEvaluationRequest):
-    """
-    Evaluate audio classification for rainforest sound detection.
-    """
-    # Get space info
-    logging.info("Chargement des données")
-    dataset = load_dataset(request.dataset_name, streaming=True, token=os.getenv("HF_TOKEN"))
-    logging.info("Données chargées")
+class evaluate_consumption(Tool):
+    name = "evaluate_consumption"
+    description = "Uses code carbon to evaluate the CO2 emissions from a given Python code"
+    inputs = {
+        "code": {
+            "type": "string",
+            "description": "The code to evaluate."
+        }
+    }
+    output_type = "string"
     
-    test_dataset = dataset["test"]
-    del dataset
+    def forward(self, code):
+        request = AudioEvaluationRequest() 
+        logging.info("Chargement des données")
+        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")
     
-    # Start tracking emissions
-    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)  
-    
-    gc.collect()
-
-    model_name = "CindyDelage/Challenge_HuggingFace_DFG_FrugalAI"
-    model = Wav2Vec2ForSequenceClassification.from_pretrained(model_name)
+        feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base")
     
-    # 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
+        test_dataset = test_dataset.map(preprocess_function, fn_kwargs={"feature_extractor": feature_extractor}, remove_columns="audio", batched=True, batch_size=32)  
+        
+        gc.collect()
     
-    classifier = pipeline("audio-classification", model=model, feature_extractor=feature_extractor, device=device)
-    predictions = []  
-    logging.info("Début des prédictions par batch")
-    i=0
-    for data in iter(test_dataset):
-        print(i)
-        if (i<=50):
-            with torch.no_grad():
-                result = classifier(np.asarray(data["input_values"]), batch_size=64)
-            predicted_label = result[0]['label']
-            label = 1 if predicted_label == 'environment' else 0
-            predictions.append(label)  
-            
-            # Nettoyer la mémoire après chaque itération
-            del result
-            del label 
-            torch.cuda.empty_cache()
-            gc.collect()
-            i=i+1
-        if(i>50):
-            break
-    logging.info("Fin des prédictions")
-    del classifier 
-    del feature_extractor 
-    
-    gc.collect()
-    # Stop tracking emissions
-    emissions_data = tracker.stop_task()
-    
-    return emissions_data
-
-request = AudioEvaluationRequest()  
-print(evaluate_audio(request))
+        model_name = "CindyDelage/Challenge_HuggingFace_DFG_FrugalAI"
+        model = Wav2Vec2ForSequenceClassification.from_pretrained(model_name)
+        
+        # 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")
+        i=0
+        for data in iter(test_dataset):
+            print(i)
+            if (i<=50):
+                with torch.no_grad():
+                    result = classifier(np.asarray(data["input_values"]), batch_size=64)
+                predicted_label = result[0]['label']
+                label = 1 if predicted_label == 'environment' else 0
+                predictions.append(label)  
+                
+                # Nettoyer la mémoire après chaque itération
+                del result
+                del label 
+                torch.cuda.empty_cache()
+                gc.collect()
+                i=i+1
+            if(i>50):
+                break
+        logging.info("Fin des prédictions")
+        del classifier 
+        del feature_extractor 
+        
+        gc.collect()
+        # Stop tracking emissions
+        emissions_data = tracker.stop_task()
+        
+        return emissions_data