Update app.py

app.py

@@ -8,14 +8,58 @@ from datasets import load_dataset, DatasetDict, Audio
 from huggingface_hub import PyTorchModelHubMixin
 import numpy as np
 
-# [Your existing code for device setup, config, SpeechInferenceDataset, SpeechClassifier]
+# 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
 
 # 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]]) 
+    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