import gradio as gr
import torch
from transformers import Speech2Text2Processor, SpeechEncoderDecoderModel
import soundfile as sf

# Load the model and processor
model = SpeechEncoderDecoderModel.from_pretrained("facebook/s2t-wav2vec2-large-en-de")
processor = Speech2Text2Processor.from_pretrained("facebook/s2t-wav2vec2-large-en-de")

# Define the transcription function
def transcribe_speech(file_info):
    # Read the audio file
    speech, _ = sf.read(file_info)
    # Process the speech
    inputs = processor(speech, sampling_rate=16_000, return_tensors="pt")
    # Generate the transcription
    generated_ids = model.generate(inputs=inputs["input_values"], attention_mask=inputs["attention_mask"])
    # Decode the generated ids to text
    transcription = processor.batch_decode(generated_ids)

    return transcription[0]

# Create the Gradio interface
iface = gr.Interface(
    fn=transcribe_speech,
    inputs=gr.Audio(label="Upload your MP3 file"),
    outputs="text",
    title="Speech to Text Conversion",
    description="Upload an MP3 file to transcribe it to text using a state-of-the-art speech-to-text model."
)


# Run the Gradio app
iface.launch()