Update app.py

app.py

@@ -8,58 +8,14 @@ from datasets import load_dataset, DatasetDict, Audio
 from huggingface_hub import PyTorchModelHubMixin
 import numpy as np
 
-# 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):
-        self.audio_data = audio_data
-        self.text_processor = text_processor
-
-    def __len__(self):
-        return len(self.audio_data)
-
-    def __getitem__(self, index):
-        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
-        decoder_input_ids = torch.tensor([[1, 1]])  # Modify as per your model's requirements
-        return input_features, decoder_input_ids
-
-# Define model class
-class SpeechClassifier(nn.Module, PyTorchModelHubMixin):
-    def __init__(self, config):
-        super(SpeechClassifier, self).__init__()
-        self.encoder = WhisperModel.from_pretrained(config["encoder"])
-        self.classifier = nn.Sequential(
-            nn.Linear(self.encoder.config.hidden_size, 4096),
-            nn.ReLU(),
-            nn.Linear(4096, 2048),
-            nn.ReLU(),
-            nn.Linear(2048, 1024),
-            nn.ReLU(),
-            nn.Linear(1024, 512),
-            nn.ReLU(),
-            nn.Linear(512, config["num_labels"])
-        )
-
-    def forward(self, input_features, decoder_input_ids):
-        outputs = self.encoder(input_features, decoder_input_ids=decoder_input_ids)
-        pooled_output = outputs['last_hidden_state'][:, 0, :]
-        logits = self.classifier(pooled_output)
-        return logits
+# [Your existing code for device setup, config, SpeechInferenceDataset, SpeechClassifier]
 
 # Prepare data function
 def prepare_data(audio_data, sampling_rate, model_checkpoint="openai/whisper-base"):
     feature_extractor = WhisperFeatureExtractor.from_pretrained(model_checkpoint)
     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
+    decoder_input_ids = torch.tensor([[1, 1]]) 
     return input_features.to(device), decoder_input_ids.to(device)
 
 # Prediction function
@@ -67,56 +23,35 @@ 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)
-    # 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
-    label = "Hypernasality Detected" if prediction == 1 else "No Hypernasality Detected"
-    return label
-
-def gradio_mic_interface(mic_input):
-    # Assuming mic_input is a tuple with the audio data and sample rate
-    if isinstance(mic_input, tuple):
-        audio_data, sample_rate = mic_input
+# Unified Gradio interface function
+def gradio_interface(audio_input):
+    if isinstance(audio_input, tuple):
+        # If the input is a tuple, it's from the microphone
+        audio_data, sample_rate = audio_input
     else:
-        # If it's not a tuple, we will try to extract data as if it's a dict
-        audio_data = mic_input['data']
-        sample_rate = mic_input['sample_rate']
+        # Otherwise, it's an uploaded file
+        with open(audio_input, "rb") as f:
+            audio_data = np.frombuffer(f.read(), np.int16)
+        sample_rate = 16000  # Assume 16kHz sample rate for uploaded files
 
     prediction = predict(audio_data, sample_rate, config)
     label = "Hypernasality Detected" if prediction == 1 else "No Hypernasality Detected"
     return label
 
-# Initialize Blocks
-demo = gr.Blocks()
-
-# Define the interfaces inside the Blocks context
-with demo:
-    mic_transcribe = gr.Interface(
-        fn=gradio_mic_interface,
-        inputs=gr.Audio(type="numpy"),  # Use numpy for real-time audio like microphone
-        outputs=gr.Textbox(label="Prediction")
-    )
-
-    file_transcribe = gr.Interface(
-        fn=gradio_file_interface,
-        inputs=gr.Audio(type="filepath"),  # Use filepath for uploaded audio files
-        outputs=gr.Textbox(label="Prediction")
-    )
-    
-    # Combine interfaces into a tabbed interface
-    gr.TabbedInterface([mic_transcribe, file_transcribe], ["Transcribe Microphone", "Transcribe Audio File"])
+# Create Gradio interface
+demo = gr.Interface(
+    fn=gradio_interface,
+    inputs=gr.Audio(type="numpy", label="Upload or Record Audio"),
+    outputs=gr.Textbox(label="Prediction")
+)
 
-# Launch the demo with debugging enabled
+# Launch the demo
 demo.launch(debug=True)