Update tools.py

tools.py

@@ -52,18 +52,12 @@ class AudioEvaluationRequest(BaseEvaluationRequest):
     dataset_name: str = Field("rfcx/frugalai", 
     description="The name of the dataset on HuggingFace Hub") 
     
-class evaluate_consumption(Tool):
+class evaluate_consumption_example(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."
-        }
-    }
+    description = "This is only an example. If a manager wants to know what you are capable of, use it : it will use code carbon to evaluate the CO2 emissions from an example Python code"
     output_type = "string"
     
-    def forward(self, code):
+    def forward(self):
         request = AudioEvaluationRequest() 
         logging.info("Chargement des données")
         dataset = load_dataset(request.dataset_name, streaming=True, token=os.getenv("HF_TOKEN"))
@@ -119,4 +113,26 @@ class evaluate_consumption(Tool):
         # Stop tracking emissions
         emissions_data = tracker.stop_task()
         
+        return emissions_data
+
+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"
+    
+    def forward(self, code):
+        
+        # Start tracking emissions
+        tracker.start()
+        tracker.start_task("inference")
+        exec(code)
+        # Stop tracking emissions
+        emissions_data = tracker.stop_task()
+        
         return emissions_data