Update app.py

app.py

@@ -1,37 +1,39 @@
 import gradio as gr
 
-from transformers import WhisperProcessor, WhisperForConditionalGeneration
-from datasets import Audio, load_dataset
-
+import torch
 import torchaudio
+from datasets import load_dataset
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
+import pandas as pd
+from sklearn.model_selection import train_test_split
+
+
+
+processor = Wav2Vec2Processor.from_pretrained("gagan3012/wav2vec2-xlsr-nepali")
+model = Wav2Vec2ForCTC.from_pretrained("gagan3012/wav2vec2-xlsr-nepali")
+
 from torchaudio.transforms import Resample
+import numpy as np
 
 
-# load model and processor
-processor = WhisperProcessor.from_pretrained("openai/whisper-medium")
-model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-medium")
-forced_decoder_ids = processor.get_decoder_prompt_ids(language="hindi", task="transcribe")
 
 
 
 def transcribe_audio(audio_file):
   input_arr, sampling_rate =torchaudio.load(audio_file)
-  input_arr = input_arr[0].numpy()
-
-  if sampling_rate != 16000:
-    resampler = Resample(orig_freq=sampling_rate, new_freq=16000)
-    input_arr = resampler(input_arr).squeeze().numpy()
-    sampling_rate = 16000
-  input_features = processor(input_arr, sampling_rate=sampling_rate, return_tensors="pt").input_features
-
-  # generate token ids
-  predicted_ids = model.generate(input_features, forced_decoder_ids=forced_decoder_ids)
-  # decode token ids to text
-  transcription = processor.batch_decode(predicted_ids)
-
-  transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)
-  print(transcription)
-  return transcription[0]
+  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]
 
 audio_input = gr.inputs.Audio(source="upload", type="filepath")