update inference code

tasks/audio.py

@@ -4,7 +4,8 @@ from datasets import load_dataset
 from sklearn.metrics import accuracy_score
 import random
 import os
-
+import tensorflow as tf
+import numpy as np
 from .utils.evaluation import AudioEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
@@ -43,6 +44,23 @@ async def evaluate_audio(request: AudioEvaluationRequest):
     # Split dataset
     train_test = dataset["train"].train_test_split(test_size=request.test_size, seed=request.test_seed)
     test_dataset = train_test["test"]
+
+
+    def compute_spectrogram(audio_array, sample_rate=16000, frame_length=256, frame_step=128):
+        spectrogram = tf.signal.stft(audio_array, frame_length=frame_length, frame_step=frame_step)
+        spectrogram = tf.abs(spectrogram)
+        return tf.expand_dims(spectrogram, axis=-1)
+
+    def preprocess(item, max_length=16000):
+        audio_array = item["audio"]["array"]
+        audio_array = tf.convert_to_tensor(audio_array, dtype=tf.float32)
+        if len(audio_array) < max_length:
+            pad_size = max_length - len(audio_array)
+            audio_array = tf.concat([audio_array, tf.zeros(pad_size)], axis=0)
+        else:
+            audio_array = audio_array[:max_length]
+        spectrogram = compute_spectrogram(audio_array)
+        return spectrogram
     
     # Start tracking emissions
     tracker.start()
@@ -52,10 +70,21 @@ async def evaluate_audio(request: AudioEvaluationRequest):
     # YOUR MODEL INFERENCE CODE HERE
     # Update the code below to replace the random baseline by your model inference within the inference pass where the energy consumption and emissions are tracked.
     #--------------------------------------------------------------------------------------------   
-    
-    # Make random predictions (placeholder for actual model inference)
+    MODEL_PATH = './model'  
+    model = tf.keras.models.load_model(MODEL_PATH)
     true_labels = test_dataset["label"]
-    predictions = [random.randint(0, 1) for _ in range(len(true_labels))]
+    predictions = []
+
+    for item in test_dataset:
+        spectrogram = preprocess(item)
+        spectrogram = tf.expand_dims(spectrogram, axis=0)  # Add batch dimension
+        pred_probs = model.predict(spectrogram, verbose=0)
+        predicted_label = np.argmax(pred_probs)
+        predictions.append(predicted_label)
+        
+    # Make random predictions (placeholder for actual model inference)
+    #true_labels = test_dataset["label"]
+    #predictions = [random.randint(0, 1) for _ in range(len(true_labels))]
     
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE