updated Model Path

app.py

@@ -1,64 +1,73 @@
-from transformers import Wav2Vec2Processor, Wav2Vec2Model
-from safetensors.torch import load_file
 import torch
-import gradio as gr
 import torchaudio
+import gradio as gr
+import os
+import time
+import numpy as np
 
-# Load processor from Hugging Face Model Hub
-processor = Wav2Vec2Processor.from_pretrained("creativepurus/accent-wav2vec2")
+from transformers import Wav2Vec2Processor, Wav2Vec2Model
+from safetensors.torch import load_file
+import torch.nn as nn
 
-# Load base model (large version)
-base_model = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-large-960h")
+# Wait for model.safetensors to be available
+model_path = "model.safetensors"
+timeout = 300  # wait up to 5 minutes
+start_time = time.time()
 
-# Define your classifier head
-classifier_head = torch.nn.Sequential(
-    torch.nn.AdaptiveAvgPool1d(1),
-    torch.nn.Flatten(),
-    torch.nn.Dropout(0.1),
-    torch.nn.Linear(1024, 2)  # 1024 hidden size for wav2vec2-large
-)
+while not os.path.exists(model_path):
+    if time.time() - start_time > timeout:
+        raise TimeoutError(f"{model_path} not found after {timeout} seconds.")
+    print(f"Waiting for {model_path} to be downloaded...")
+    time.sleep(5)
+
+# Load processor
+processor = Wav2Vec2Processor.from_pretrained("creativepurus/accent-wav2vec2")
 
-# Load fine-tuned classifier weights
-state_dict = load_file("model.safetensors", device="cpu")
-classifier_head.load_state_dict(state_dict)
+# Load model weights from model.safetensors
+state_dict = load_file(model_path, device="cpu")
 
-# Combine base model + classifier head
-class AccentClassifier(torch.nn.Module):
-    def __init__(self, base, head):
-        super().__init__()
-        self.base = base
-        self.head = head
+# Define the same model architecture used during training
+class Wav2Vec2Classifier(nn.Module):
+    def __init__(self):
+        super(Wav2Vec2Classifier, self).__init__()
+        self.wav2vec2 = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-large-960h")
+        self.dropout = nn.Dropout(0.3)
+        self.classifier = nn.Linear(self.wav2vec2.config.hidden_size, 2)
 
-    def forward(self, input_values):
-        with torch.no_grad():
-            features = self.base(input_values).last_hidden_state
-        logits = self.head(features.transpose(1, 2))
+    def forward(self, input_values, attention_mask=None):
+        outputs = self.wav2vec2(input_values, attention_mask=attention_mask)
+        hidden_states = outputs.last_hidden_state
+        pooled = hidden_states.mean(dim=1)
+        pooled = self.dropout(pooled)
+        logits = self.classifier(pooled)
         return logits
 
-model = AccentClassifier(base_model, classifier_head)
+# Instantiate and load the model
+model = Wav2Vec2Classifier()
+model.load_state_dict(state_dict)
 model.eval()
 
-# Inference function
-def predict(audio_path):
-    # Load and preprocess audio
-    waveform, sample_rate = torchaudio.load(audio_path)
+# Prediction function
+def predict_accent(audio):
+    waveform, sample_rate = torchaudio.load(audio)
     if sample_rate != 16000:
         resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)
         waveform = resampler(waveform)
-    
-    inputs = processor(waveform.squeeze().numpy(), sampling_rate=16000, return_tensors="pt", padding=True)
-    logits = model(inputs.input_values)
-    probs = torch.nn.functional.softmax(logits, dim=1)
-    
-    labels = ["Canadian English", "England English"]
-    return {labels[i]: float(probs[0][i]) for i in range(2)}
+    input_values = processor(waveform.squeeze().numpy(), return_tensors="pt", sampling_rate=16000).input_values
+    with torch.no_grad():
+        logits = model(input_values)
+        predicted_class_id = logits.argmax().item()
+    label_map = {0: "Canadian English", 1: "England English"}
+    return label_map[predicted_class_id]
 
-# Gradio Interface
+# Gradio UI
 interface = gr.Interface(
-    fn=predict,
-    inputs=gr.Audio(source="upload", type="filepath"),
-    outputs=gr.Label(num_top_classes=2),
-    title="Accent Classification with Wav2Vec2-Large"
+    fn=predict_accent,
+    inputs=gr.Audio(source="upload", type="filepath", label="Upload or Record Audio (WAV)"),
+    outputs=gr.Textbox(label="Predicted Accent"),
+    title="Accent Classification",
+    description="This app classifies English accents as either Canadian or England using a fine-tuned Wav2Vec2 model.",
+    allow_flagging="never"
 )
 
 interface.launch()