Update app.py

app.py

@@ -28,7 +28,7 @@ class SpeechInferenceDataset(Dataset):
         inputs = self.text_processor(self.audio_data[index]["audio"]["array"],
                                      return_tensors="pt",
                                      sampling_rate=self.audio_data[index]["audio"]["sampling_rate"])
-        input_features = inputs.input_features
+        input_features = inputs.input_features.squeeze(0)
         decoder_input_ids = torch.tensor([[1, 1]])  # Modify as per your model's requirements
         return input_features, decoder_input_ids
 
@@ -58,6 +58,9 @@ class SpeechClassifier(nn.Module, PyTorchModelHubMixin):
 # Prepare data function
 def prepare_data(audio_data, sampling_rate, model_checkpoint="openai/whisper-base"):
 
+    # Convert audio data to float32
+    audio_data = audio_data.astype(np.float32)
+    
     # Resample audio data to 16000 Hz
     audio_data_resampled = librosa.resample(audio_data, orig_sr=sampling_rate, target_sr=16000)
 
@@ -68,12 +71,15 @@ def prepare_data(audio_data, sampling_rate, model_checkpoint="openai/whisper-bas
     dataset = SpeechInferenceDataset([{"audio": {"array": audio_data_resampled, "sampling_rate": 16000}}],
                                      text_processor=feature_extractor)
 
-    return dataset
+    dataloader = DataLoader(dataset, batch_size=1)
+
+    return dataloader
+    # return dataset
 
 
 # Prediction function
 def predict(audio_data, sampling_rate, config):
-    input_features, decoder_input_ids = prepare_data(audio_data, sampling_rate, config["encoder"])
+    dataloader = prepare_data(audio_data, sampling_rate, config["encoder"])
 
     model = SpeechClassifier(config).to(device)
     # Here we load the model from Hugging Face Hub
@@ -81,23 +87,43 @@ def predict(audio_data, sampling_rate, config):
 
     model.eval()
     with torch.no_grad():
-        logits = model(input_features, decoder_input_ids)
-        predicted_ids = int(torch.argmax(logits, dim=-1))
+        for input_features, decoder_input_ids in dataloader:
+            input_features = input_features.to(device)
+            decoder_input_ids = decoder_input_ids.to(device)
+            logits = model(input_features, decoder_input_ids)
+            predicted_ids = int(torch.argmax(logits, dim=-1))
     return predicted_ids
+    
+    # input_features, decoder_input_ids = prepare_data(audio_data, sampling_rate, config["encoder"])
+
+    # model = SpeechClassifier(config).to(device)
+    # # Here we load the model from Hugging Face Hub
+    # model.load_state_dict(torch.hub.load_state_dict_from_url("https://huggingface.co/jcho02/whisper_cleft/resolve/main/pytorch_model.bin", map_location=device))
+
+    # model.eval()
+    # with torch.no_grad():
+    #     logits = model(input_features, decoder_input_ids)
+    #     predicted_ids = int(torch.argmax(logits, dim=-1))
+    # return predicted_ids
 
 # Gradio Interface functions
 def gradio_file_interface(uploaded_file):
     # Assuming the uploaded_file is a filepath (str)
-    with open(uploaded_file, "rb") as f:
-        audio_data = np.frombuffer(f.read(), np.int16)
-    prediction = predict(audio_data, 16000, config)  # Assume 16kHz sample rate
+    audio_data, sampling_rate = librosa.load(uploaded_file, sr=None)
+    prediction = predict(audio_data, sampling_rate, config)
     label = "Hypernasality Detected" if prediction == 1 else "No Hypernasality Detected"
     return label
 
+    # with open(uploaded_file, "rb") as f:
+    #     audio_data = np.frombuffer(f.read(), np.int16)
+    # prediction = predict(audio_data, 16000, config)  # Assume 16kHz sample rate
+    # label = "Hypernasality Detected" if prediction == 1 else "No Hypernasality Detected"
+    # return label
+
 def gradio_mic_interface(mic_input):
     # mic_input is a tuple with sample_rate and data as entries
     # (44100, array([   0,    0,    0, ..., -153, -140, -120], dtype=int16))
-    prediction = predict(mic_input[1], mic_input[0], config)
+    prediction = predict(mic_input[1].astype(np.float32), mic_input[0], config)
     label = "Hypernasality Detected" if prediction == 1 else "No Hypernasality Detected"
     return label
 
@@ -119,5 +145,5 @@ with gr.Blocks() as demo:
             outputs=gr.Textbox(label="Prediction")
         )
 
-# Launch the demo
-demo.launch()
+# Launch the demo with debugging enabled
+demo.launch(debug=True)