Update app.py

app.py

@@ -1,51 +1,19 @@
-import gradio as gr
-from transformers import Wav2Vec2Processor, Wav2Vec2BertForCTC
+import soundfile as sf
 import torch
-import librosa
-import numpy as np
-
-# Load the correct processor and model
-model_id = "kdcyberdude/w2v-bert-punjabi"
-processor = Wav2Vec2Processor.from_pretrained(model_id)
-model = Wav2Vec2BertForCTC.from_pretrained(model_id)
-
-def transcribe_audio(audio_file):
-    try:
-        # Load and preprocess the audio
-        audio, rate = librosa.load(audio_file, sr=16000)  # Resample to 16 kHz
-        if len(audio.shape) > 1:  # If stereo, convert to mono
-            audio = np.mean(audio, axis=1)
-        
-        # Normalize audio to match expected input range [-1, 1]
-        audio = librosa.util.normalize(audio)
-        
-        # Split into manageable chunks (30 seconds each)
-        chunk_size = int(30 * rate)  # 30 seconds in samples
-        transcription = []
-
-        for i in range(0, len(audio), chunk_size):
-            chunk = audio[i:i + chunk_size]
-            input_values = processor(chunk, sampling_rate=16000, return_tensors="pt").input_values
-            
-            # Perform inference
-            with torch.no_grad():
-                logits = model(input_values).logits
-            
-            # Decode predicted IDs to text
-            predicted_ids = torch.argmax(logits, dim=-1)
-            transcription.append(processor.batch_decode(predicted_ids)[0])
-
-        return " ".join(transcription)
-    except Exception as e:
-        return f"Error: {str(e)}"
-
-# Gradio interface setup
-iface = gr.Interface(
-    fn=transcribe_audio,
-    inputs=gr.Audio(type="filepath"),
-    outputs=gr.Textbox(label="Punjabi Transcription"),
-    title="Punjabi Audio Transcription",
-    description="Upload an audio file to transcribe Punjabi speech."
-)
-
-iface.launch()
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
+import argparse
+def parse_transcription(wav_file):
+    # load pretrained model
+    processor = Wav2Vec2Processor.from_pretrained("addy88/wav2vec2-punjabi-stt")
+    model = Wav2Vec2ForCTC.from_pretrained("addy88/wav2vec2-punjabi-stt")
+    # load audio
+    audio_input, sample_rate = sf.read(wav_file)
+    # pad input values and return pt tensor
+    input_values = processor(audio_input, sampling_rate=sample_rate, return_tensors="pt").input_values
+    # INFERENCE
+    # retrieve logits & take argmax
+    logits = model(input_values).logits
+    predicted_ids = torch.argmax(logits, dim=-1)
+    # transcribe
+    transcription = processor.decode(predicted_ids[0], skip_special_tokens=True)
+    print(transcription)