Update app.py

app.py

@@ -7,6 +7,9 @@ from transformers import WhisperModel, WhisperFeatureExtractor
 import datasets
 from datasets import load_dataset, DatasetDict, Audio
 from huggingface_hub import PyTorchModelHubMixin
+import numpy as np
+import tempfile
+import os
 
 # Ensure you have the device setup (cuda or cpu)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -25,7 +28,8 @@ class SpeechInferenceDataset(Dataset):
                                      return_tensors="pt",
                                      sampling_rate=self.audio_data[index]["audio"]["sampling_rate"])
         input_features = inputs.input_features
-        decoder_input_ids = torch.tensor([[1, 1]])  # Modify as per your model's requirements
+        # Modify decoder_input_ids as per your model's requirements
+        decoder_input_ids = torch.tensor([[1, 1]])
         return input_features, decoder_input_ids
 
 # Define model class
@@ -51,44 +55,34 @@ class SpeechClassifier(nn.Module, PyTorchModelHubMixin):
         logits = self.classifier(pooled_output)
         return logits
 
-# Prepare data function
-def prepare_data(audio_file_path, model_checkpoint="openai/whisper-base"):
+# Prepare data function (may need to update for numpy input)
+def prepare_data(audio_data, sampling_rate, model_checkpoint="openai/whisper-base"):
     feature_extractor = WhisperFeatureExtractor.from_pretrained(model_checkpoint)
-    inference_data = datasets.Dataset.from_dict({"path": [audio_file_path], "audio": [audio_file_path]}).cast_column("audio", Audio(sampling_rate=16_000))
-    inference_dataset = SpeechInferenceDataset(inference_data, feature_extractor)
-    inference_loader = DataLoader(inference_dataset, batch_size=1, shuffle=False)
-    input_features, decoder_input_ids = next(iter(inference_loader))
-    input_features = input_features.squeeze(1).to(device)
-    decoder_input_ids = decoder_input_ids.squeeze(1).to(device)
-    return input_features, decoder_input_ids
+    # ... your logic for preparing data ...
+    # Must return tensor that your model can process
+    pass
 
-# Prediction function
-def predict(audio_file_path, config={"encoder": "openai/whisper-base", "num_labels": 2}):
-    input_features, decoder_input_ids = prepare_data(audio_file_path)
-    
-    # Load the model from Hugging Face Hub
-    model = SpeechClassifier(config)
-    model.to(device)
-    # Use the correct method to load your model (this is an example and may not directly apply)
-    model.load_state_dict(torch.load(model.push_from_hub("jcho02/whisper_cleft")))
-    model.eval()
-
-    with torch.no_grad():
-        logits = model(input_features, decoder_input_ids)
-        predicted_ids = int(torch.argmax(logits, dim=-1))
-    return predicted_ids
+# Prediction function (may need to update for numpy input)
+def predict(audio_data, sampling_rate, config={"encoder": "openai/whisper-base", "num_labels": 2}):
+    # Load the model from Hugging Face Hub (ensure correct loading mechanism)
+    model = SpeechClassifier(config).to(device)
+    # ... your logic for prediction using model ...
+    # Must return a prediction
+    pass
 
 # Gradio Interface function for uploaded files
 def gradio_file_interface(uploaded_file):
-    with open(uploaded_file.name, "wb") as f:
-        f.write(uploaded_file.read())
-    prediction = predict(uploaded_file.name)
+    # Gradio passes a file path as a string for uploaded files
+    prediction = predict(uploaded_file, config)
     label = "Hypernasality Detected" if prediction == 1 else "No Hypernasality Detected"
     return label
 
 # Gradio Interface function for microphone input
 def gradio_mic_interface(mic_input):
-    prediction = predict(mic_input.name)
+    # Gradio passes mic input as a numpy array and sample rate
+    audio_data = mic_input['data']
+    sampling_rate = mic_input['sample_rate']
+    prediction = predict(audio_data, sampling_rate, config)
     label = "Hypernasality Detected" if prediction == 1 else "No Hypernasality Detected"
     return label
 
@@ -99,18 +93,18 @@ demo = gr.Blocks()
 with demo:
     mic_transcribe = gr.Interface(
         fn=gradio_mic_interface,
-        inputs=gr.Audio(),  # No type needed for microphone input
+        inputs=gr.Audio(type="numpy"),  # Receives numpy array for mic input
         outputs=gr.Textbox(label="Prediction")
     )
 
     file_transcribe = gr.Interface(
         fn=gradio_file_interface,
-        inputs=gr.Audio(type="filepath"),  # Specify filepath for file upload
+        inputs=gr.Audio(type="filepath"),  # Receives file path for file upload
         outputs=gr.Textbox(label="Prediction")
     )
     
     # Use a tabbed interface to switch between the microphone and file upload interfaces
     gr.TabbedInterface([mic_transcribe, file_transcribe], ["Transcribe Microphone", "Transcribe Audio File"])
 
-# Launch the demo with debugging enabled to catch any potential errors early on
+# Launch the demo with debugging enabled
 demo.launch(debug=True)