import gradio as gr
import librosa
from transformers import AutoFeatureExtractor, AutoTokenizer, SpeechEncoderDecoderModel

model_name = "facebook/wav2vec2-xls-r-2b-21-to-en",
    
feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
model = SpeechEncoderDecoderModel.from_pretrained(model_name)

def process_audio_file(file):
    data, sr = librosa.load(file)
    if sr != 16000:
        data = librosa.resample(data, sr, 16000)
    input_values = feature_extractor(data, return_tensors="pt").input_values
    return input_values
    
def transcribe(file):
    input_values = process_audio_file(file)
    
    sequences = model.generate(input_values, num_beams=1, max_length=30)
    
    transcription = tokenizer.batch_decode(sequences, skip_special_tokens=True)
    return transcription[0]

iface = gr.Interface(
    fn=transcribe, 
    inputs=[
        gr.inputs.Audio(source="microphone", type='filepath'),
    ],
    outputs="text",
    layout="horizontal",
    theme="huggingface",
    title="XLS-R 2B 21-to-EN Speech Translation",
    description="A simple interface to translate from 21 spoken languages to written English.",
)
iface.launch()