Update app.py

app.py

@@ -1,4 +1,3 @@
-# Import necessary libraries
 import gradio as gr
 import torch
 import torch.nn as nn
@@ -8,12 +7,13 @@ 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")
 
+# Define the config for your model
+config = {"encoder": "openai/whisper-base", "num_labels": 2}
+
 # Define data class
 class SpeechInferenceDataset(Dataset):
     def __init__(self, audio_data, text_processor):
@@ -28,8 +28,7 @@ class SpeechInferenceDataset(Dataset):
                                      return_tensors="pt",
                                      sampling_rate=self.audio_data[index]["audio"]["sampling_rate"])
         input_features = inputs.input_features
-        # Modify decoder_input_ids as per your model's requirements
-        decoder_input_ids = torch.tensor([[1, 1]])
+        decoder_input_ids = torch.tensor([[1, 1]])  # Modify as per your model's requirements
         return input_features, decoder_input_ids
 
 # Define model class
@@ -55,34 +54,40 @@ class SpeechClassifier(nn.Module, PyTorchModelHubMixin):
         logits = self.classifier(pooled_output)
         return logits
 
-# Prepare data function (may need to update for numpy input)
+# Prepare data function
 def prepare_data(audio_data, sampling_rate, model_checkpoint="openai/whisper-base"):
     feature_extractor = WhisperFeatureExtractor.from_pretrained(model_checkpoint)
-    # ... your logic for preparing data ...
-    # Must return tensor that your model can process
-    pass
+    inputs = feature_extractor(audio_data, sampling_rate=sampling_rate, return_tensors="pt")
+    input_features = inputs.input_features
+    decoder_input_ids = torch.tensor([[1, 1]])  # Modify as per your model's requirements
+    return input_features.to(device), decoder_input_ids.to(device)
 
-# 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)
+# Prediction function
+def predict(audio_data, sampling_rate, config):
+    input_features, decoder_input_ids = prepare_data(audio_data, sampling_rate, config["encoder"])
+    
     model = SpeechClassifier(config).to(device)
-    # ... your logic for prediction using model ...
-    # Must return a prediction
-    pass
+    # 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 function for uploaded files
+# Gradio Interface functions
 def gradio_file_interface(uploaded_file):
-    # Gradio passes a file path as a string for uploaded files
-    prediction = predict(uploaded_file, config)
+    # 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
     label = "Hypernasality Detected" if prediction == 1 else "No Hypernasality Detected"
     return label
 
-# Gradio Interface function for microphone input
 def gradio_mic_interface(mic_input):
-    # 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)
+    # mic_input is a dictionary with 'data' and 'sample_rate' keys
+    prediction = predict(mic_input['data'], mic_input['sample_rate'], config)
     label = "Hypernasality Detected" if prediction == 1 else "No Hypernasality Detected"
     return label
 
@@ -93,18 +98,17 @@ demo = gr.Blocks()
 with demo:
     mic_transcribe = gr.Interface(
         fn=gradio_mic_interface,
-        inputs=gr.Audio(type="numpy"),  # Receives numpy array for mic input
+        inputs=gr.Audio(source="microphone", type="numpy"),  # Correct type for microphone
         outputs=gr.Textbox(label="Prediction")
     )
 
     file_transcribe = gr.Interface(
         fn=gradio_file_interface,
-        inputs=gr.Audio(type="filepath"),  # Receives file path for file upload
+        inputs=gr.Audio(source="upload", type="file"),  # Correct type 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
+# Launch the demo
 demo.launch(debug=True)