Update app.py

app.py

@@ -1,46 +1,48 @@
 import torch
-import gradio as gr
 import torchaudio
-from transformers import Wav2Vec2ForCTC, AutoProcessor
-from quanto import qint8, quantize, freeze
+import gradio as gr
+from transformers import Wav2Vec2BertProcessor, Wav2Vec2BertForCTC
+
+# Set device
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load and quantize the model
-model_name = "cdactvm/w2v-bert-punjabi"
-model = Wav2Vec2ForCTC.from_pretrained(model_name)  # Ensure it's a CTC model
-processor = AutoProcessor.from_pretrained(model_name)
+# Load processor & model
+model_name = "cdactvm/w2v-bert-punjabi"  # Change if using a Punjabi ASR model
+processor = Wav2Vec2BertProcessor.from_pretrained(model_name)
+model = Wav2Vec2BertForCTC.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(device)
 
-# Quantization
-quantize(model, weights=qint8, activations=None)
-freeze(model)
+def transcribe(audio_path):
+    # Load audio file
+    waveform, sample_rate = torchaudio.load(audio_path)
 
-# Audio transcription function
-def transcribe(audio):
-    waveform, sample_rate = torchaudio.load(audio)
+    # Convert stereo to mono (if needed)
+    if waveform.shape[0] > 1:
+        waveform = torch.mean(waveform, dim=0, keepdim=True)
 
-    # Ensure 16kHz sample rate
+    # Resample to 16kHz
     if sample_rate != 16000:
-        waveform = torchaudio.transforms.Resample(sample_rate, 16000)(waveform)
+        waveform = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)(waveform)
 
     # Process audio
     inputs = processor(waveform.squeeze(0), sampling_rate=16000, return_tensors="pt")
+    inputs = {key: val.to(device, dtype=torch.bfloat16) for key, val in inputs.items()}
 
-    # Run inference
+    # Get logits & transcribe
     with torch.no_grad():
-        logits = model(**inputs).logits  # Ensure model has 'logits'
-
-    # Decode transcription
+        logits = model(**inputs).logits
     predicted_ids = torch.argmax(logits, dim=-1)
     transcription = processor.batch_decode(predicted_ids)[0]
-    
+
     return transcription
 
-# Gradio UI
-iface = gr.Interface(
+# Gradio Interface
+app = gr.Interface(
     fn=transcribe,
-    inputs=gr.Audio(sources="upload", type="filepath"),
+    inputs=gr.Audio(source="upload", type="filepath"),
     outputs="text",
-    title="Punjabi Speech Recognition",
-    description="Upload an audio file and get a Punjabi transcription using a quantized model.",
+    title="Punjabi Speech-to-Text",
+    description="Upload an audio file and get the transcription in Punjabi."
 )
 
-iface.launch()
+if __name__ == "__main__":
+    app.launch()