Update app.py

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
 from PIL import Image
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, Wav2Vec2Processor, Wav2Vec2ForCTC
 import soundfile as sf
 import torch
-from speech_recognition import AudioFile, Recognizer
+import numpy as np
 
 model_name_or_path = "microsoft/DialoGPT-large"
 
@@ -17,7 +17,9 @@ model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True,
 )
 
-
+# Initialize the Wav2Vec2 model and processor
+wav2vec2_processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
+wav2vec2_model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
 
 # Function to handle text input
 def handle_text(text):
@@ -32,36 +34,24 @@ def handle_image(img):
     return "This image seems nice!"
 
 # Function to handle audio input
-def stt(audio: object, language: str) -> str:
-    """Converts speech to text.
-
-    Args:
-        audio: record of user speech
-
-    Returns:
-        text (str): recognized speech of user
-    """
-    
-    # Create a Recognizer object
-    r = Recognizer()
-    # Open the audio file
-    with AudioFile(audio) as source:
-        # Listen for the data (load audio to memory)
-        audio_data = r.record(source)
-        # Transcribe the audio using Google's speech-to-text API
-        text = r.recognize_google(audio_data, language=language)
-    return text
+def handle_audio(audio):
+    # gradio Audio returns a tuple (sample_rate, audio_np_array)
+    # we only need the audio data, hence accessing the second element
+    audio = audio[1]
+    input_values = wav2vec2_processor(audio, sampling_rate=16_000, return_tensors="pt").input_values
+    logits = wav2vec2_model(input_values).logits
+    predicted_ids = torch.argmax(logits, dim=-1)
+    transcriptions = wav2vec2_processor.decode(predicted_ids[0])
+    return handle_text(transcriptions)
 
 def chatbot(text, img, audio):
     text_output = handle_text(text) if text is not None else ''
     img_output = handle_image(img) if img is not None else ''
-    audio_output = handle_text(stt(audio,'english')) if audio is not None else ''
+    audio_output = handle_audio(audio) if audio is not None else ''
 
     outputs = [o for o in [text_output, img_output, audio_output] if o]
     return "\n".join(outputs)
 
-
-
 iface = gr.Interface(
     fn=chatbot,
     inputs=[