Update app.py

app.py

@@ -4,53 +4,31 @@ subprocess.run(["pip", "install", "transformers"])
 subprocess.run(["pip", "install", "torchaudio", "--upgrade"])
 
 import gradio as gr
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
-import torchaudio
-import torch
-
-# Load model and processor
-processor = Wav2Vec2Processor.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-italian")
-model = Wav2Vec2ForCTC.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-italian")
-
-def preprocess_audio(audio_data):
-    # Apply any custom preprocessing to the audio data here if needed
-    return processor(audio_data, return_tensors="pt").input_features
-
-# Function to perform ASR on audio data
-def transcribe_audio(input_features):
-    print("Received audio data:", input_features)  # Debug print
-
-    # Check if audio_data is None or not a tuple of length 2
-    if audio_data is None or not isinstance(input_features, tuple) or len(input_features) != 2:
-        return "Invalid audio data format."
-
-    sample_rate, waveform = input_features
-
-    # Check if waveform is None or not a NumPy array
-    if waveform is None or not isinstance(waveform, torch.Tensor):
-        return "Invalid audio data format."
-
-    try:
-        # Convert audio data to mono and normalize
-        audio_data = torchaudio.transforms.Resample(sample_rate, 100000)(waveform)
-        audio_data = torchaudio.functional.gain(input_features, gain_db=5.0)
-
-        # Apply custom preprocessing to the audio data if needed
-        input_values = processor(input_features[0], return_tensors="pt").input_values
-
-        # Perform ASR
-        with torch.no_grad():
-            logits = model(input_values).logits
-
-        # Decode the output
-        predicted_ids = torch.argmax(logits, dim=-1)
-        transcription = processor.batch_decode(predicted_ids)
-
-        return transcription[0]
-
-    except Exception as e:
-        return f"An error occurred: {str(e)}"
-
-# Create Gradio interface
-audio_input = gr.Audio(sources=["microphone"])
-gr.Interface(fn=transcribe_audio, inputs=audio_input, outputs="text").launch()
+from transformers import WhisperProcessor, WhisperForConditionalGeneration
+
+# Load Whisper ASR model and processor
+model_name = "openai/whisper-small"
+processor = WhisperProcessor.from_pretrained(model_name)
+model = WhisperForConditionalGeneration.from_pretrained(model_name)
+forced_decoder_ids = processor.get_decoder_prompt_ids(language="italian", task="transcribe")
+
+def transcribe_audio(input_audio):
+    # Process audio using the Whisper processor
+    input_features = processor(input_audio, 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, skip_special_tokens=True)
+    
+    return transcription[0]
+
+iface = gr.Interface(
+    fn=transcribe_audio,
+    inputs=gr.Audio(source="microphone", type="wav", label="Speak"),
+    outputs="text",
+    live=True
+)
+
+iface.launch()