Update app.py

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
-
+import librosa
 import torch
 import torchaudio
 from datasets import load_dataset
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, pipeline
+from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
 import pandas as pd
 from sklearn.model_selection import train_test_split
 
@@ -13,41 +13,17 @@ from pydub import AudioSegment
 
 
 
-# processor = Wav2Vec2Processor.from_pretrained("RikeshSilwal/wav2vec2-nepali")
-# model = Wav2Vec2ForCTC.from_pretrained("RikeshSilwal/wav2vec2-nepali")
-
-processor = Wav2Vec2Processor.from_pretrained("RikeshSilwal/wav2vec2-nepali-rikeshsilwal")
-model = Wav2Vec2ForCTC.from_pretrained("RikeshSilwal/wav2vec2-nepali-rikeshsilwal")
-
 from torchaudio.transforms import Resample
 import numpy as np
 
 
 
-
-
-# def transcribe_audio(audio_file):
-#   input_arr, sampling_rate =torchaudio.load(audio_file)
-#   resampler = Resample(orig_freq=sampling_rate, new_freq=16000)
-#   input_arr = resampler(input_arr).squeeze().numpy()
-#   sampling_rate = 16000
-#   inputs = processor(input_arr, sampling_rate=16_000, return_tensors="pt", padding=True)
-
-#   with torch.no_grad():
-#     logits = model(inputs.input_values, attention_mask=inputs.attention_mask).logits
-
-#   predicted_ids = torch.argmax(logits, dim=-1)
-
-#   predicted_words=  processor.batch_decode(predicted_ids)
-
-#   return predicted_words[0]
-
 def transcribe_audio(audio_file):
     audio = AudioSegment.from_wav(audio_file)
 
     device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 
-    input_arr, sampling_rate =torchaudio.load(audio_file)
+    input_arr, sampling_rate =librosa.load(audio_file)
     # Create TorchGating instance
     tg = TG(sr=sampling_rate, nonstationary=True).to(device)
     try:
@@ -55,13 +31,39 @@ def transcribe_audio(audio_file):
     except:
       input_arr = input_arr
     if sampling_rate != 16000:
-      resampler = torchaudio.transforms.Resample(orig_freq=sampling_rate, new_freq=16000)
-      input_arr = resampler(input_arr).squeeze().numpy()
+      input_arr = librosa.resample(input_arr, orig_sr=sampling_rate, target_sr=16000)
+
+    
+    MODEL_NAME = "rikeshsilwalekg/whisper-small-wer35-ekg"
+
+    torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
+    model = AutoModelForSpeechSeq2Seq.from_pretrained(
+        MODEL_NAME, torch_dtype=torch_dtype, use_safetensors=True
+    )
+    model.to(device)
 
-    recognizer = pipeline("automatic-speech-recognition", model="Harveenchadha/vakyansh-wav2vec2-nepali-nem-130")
-    prediction = recognizer(input_arr, chunk_length_s=5, stride_length_s=(2,1))
-    prediction = recognizer(input_arr)
-    prediction = re.sub('[<s>]' , '' , str(prediction['text']))
+    processor = AutoProcessor.from_pretrained(MODEL_NAME)
+
+    pipe = pipeline(
+        "automatic-speech-recognition",
+        model=model,
+        tokenizer=processor.tokenizer,
+        feature_extractor=processor.feature_extractor,
+        max_new_tokens=128,
+        chunk_length_s=30,
+        batch_size=16,
+        return_timestamps=False,
+        torch_dtype=torch_dtype,
+        device=device,
+    )
+
+    # return_timestamps=True for sentence level timestaps
+    # for word level timestamps return_timestamps="word"
+    prediction = pipe(input_arr)
+
+    
+    prediction = prediction['text']
 
 audio_input = gr.inputs.Audio(source="upload", type="filepath")
 
@@ -71,4 +73,3 @@ iface = gr.Interface(fn=transcribe_audio, inputs=audio_input,
                              button")
 
 iface.launch(inline=False)
-