Update app.py

app.py

@@ -1,16 +1,14 @@
 import subprocess
-subprocess.run(["pip", "install", "datasets"])
 subprocess.run(["pip", "install", "transformers"])
-subprocess.run(["pip", "install", "torch", "torchvision", "torchaudio", "-f", "https://download.pytorch.org/whl/torch_stable.html"])
+subprocess.run(["pip", "install", "torchaudio"])
 
 import gradio as gr
-from transformers import WhisperProcessor, WhisperForConditionalGeneration
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
 import torchaudio
 
 # Load model and processor
-processor = WhisperProcessor.from_pretrained("openai/whisper-large")
-model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-large")
-model.config.forced_decoder_ids = None
+processor = Wav2Vec2Processor.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-italian")
+model = Wav2Vec2ForCTC.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-italian")
 
 # Function to perform ASR on audio data
 def transcribe_audio(audio_data):
@@ -18,18 +16,20 @@ def transcribe_audio(audio_data):
     audio_data = torchaudio.functional.to_mono(audio_data)
     audio_data = torchaudio.functional.gain(audio_data, gain_db=5.0)
 
-    # Resample if needed (Whisper model requires 16 kHz sampling rate)
+    # Resample if needed (Wav2Vec2 model requires 16 kHz sampling rate)
     if audio_data[1] != 16000:
         audio_data = torchaudio.transforms.Resample(audio_data[1], 16000)(audio_data[0])
 
     # Apply custom preprocessing to the audio data if needed
-    processed_input = processor(audio_data[0].numpy(), return_tensors="pt").input_features
+    input_values = processor(audio_data[0].numpy(), return_tensors="pt").input_values
 
-    # Generate token ids
-    predicted_ids = model.generate(processed_input)
+    # Perform ASR
+    with torch.no_grad():
+        logits = model(input_values).logits
 
-    # Decode token ids to text
-    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)
+    # Decode the output
+    predicted_ids = torch.argmax(logits, dim=-1)
+    transcription = processor.batch_decode(predicted_ids)
 
     return transcription[0]