Update app.py

app.py

@@ -26,7 +26,7 @@ def is_white_noise(audio, threshold=0.75):
     kurt = kurtosis(audio)
     return np.abs(kurt) < 0.1 and np.mean(np.abs(audio)) < threshold
 
-def create_audio_dataframe(audio_tuple, target_sr=16000, target_duration=1.0):
+def process_audio(audio_tuple, target_sr=16000, target_duration=1.0):
     data = []
     target_length = int(target_sr * target_duration)
 
@@ -49,60 +49,41 @@ def create_audio_dataframe(audio_tuple, target_sr=16000, target_duration=1.0):
     if len(audio_data) > target_length:
         for i in range(0, len(audio_data), target_length):
             segment = audio_data[i:i + target_length]
-            if len(segment) == target_length:
-                if not is_white_noise(segment):
-                    data.append({"audio": segment})
+            if len(segment) == target_length and not is_white_noise(segment):
+                data.append(segment)
     else:
         if not is_white_noise(audio_data):
-            data.append({"audio": audio_data})
+            data.append(audio_data)
 
-    df = pd.DataFrame(data)
-    return df
+    return data
 
-def convert_bytes_to_float64(byte_list):
-    return [float(i) for i in byte_list]
-
-def preprocess_function(examples):
-    audio_arrays = examples["audio"]
+def preprocess_audio(audio_segments):
     inputs = processor(
-        audio_arrays,
+        audio_segments,
         padding=True,
         sampling_rate=processor.sampling_rate,
         max_length=int(processor.sampling_rate * 1),
         truncation=True,
+        return_tensors="pt"  # Directamente retorna tensores de PyTorch
     )
     return inputs
 
 def predict_audio(audio):
-    df = create_audio_dataframe(audio)
-    df['audio'] = df['audio'].apply(convert_bytes_to_float64)
+    # Procesar el audio y obtener las listas de numpy
+    audio_segments = process_audio(audio)
     
-    # Convertir el dataframe a Dataset
-    predict_dataset = Dataset.from_pandas(df)
-    dataset = DatasetDict({
-        'train': predict_dataset
-    })
-
-    if '__index_level_0__' in dataset['train'].column_names:
-        dataset['train'] = dataset['train'].remove_columns(['__index_level_0__'])
-
-    encoded_dataset = dataset.map(preprocess_function, remove_columns=["audio"], batched=True)
-
-    # Crear el Trainer para la predicción
-    trainer = Trainer(
-        model=model,  
-        eval_dataset=encoded_dataset["train"]
-    )
+    # Preprocesar el audio (aplica directamente al array numpy)
+    inputs = preprocess_audio(audio_segments)
 
     # Realizar las predicciones
-    predictions_output = trainer.predict(encoded_dataset["train"].with_format("torch"))
+    with torch.no_grad():
+        outputs = model(**inputs)
 
-    # Obtener las predicciones y etiquetas verdaderas
-    predictions = predictions_output.predictions
-    labels = predictions_output.label_ids
+    # Obtener los logits de las predicciones
+    logits = outputs.logits
 
     # Convertir logits a probabilidades
-    probabilities = F.softmax(torch.tensor(predictions), dim=-1).numpy()
+    probabilities = torch.nn.functional.softmax(logits, dim=-1).numpy()
     predicted_classes = probabilities.argmax(axis=1)
 
     # Obtener la etiqueta más común