import gradio as gr
from transformers import pipeline
from librosa import resample
import numpy as np

def transcribe(input_audio, model_id, denormalize):
    pipe = pipeline(
        "automatic-speech-recognition",
        model=model_id,
        device="cpu"
    )
    sr, speech = input_audio
    # Convert to mono if stereo
    if speech.ndim > 1:
        speech = speech.mean(axis=1)
    # Convert to float32 if needed
    if speech.dtype != "float32":
        speech = speech.astype(np.float32)
    # Resample if sampling rate is not 16kHz
    if sr!=16000:
        speech = resample(speech, orig_sr=sr, target_sr=16000)
    output = pipe(speech, chunk_length_s=30, stride_length_s=5)['text']
    # Optional: restore punctuation and capitalization
    if denormalize:
        from punctuators.models import PunctCapSegModelONNX
        DENORM_MODEL_ID = "1-800-BAD-CODE/xlm-roberta_punctuation_fullstop_truecase"
        denorm_model = PunctCapSegModelONNX.from_pretrained(
            DENORM_MODEL_ID, ort_providers=["CPUExecutionProvider"]
        )
        output = denorm_model.infer(
                texts=[output],
                apply_sbd=False,
                batch_size_tokens=4096,
                num_workers=0,
            )[0]
    return output

model_ids_list = [
 "GetmanY1/wav2vec2-base-fi-150k-finetuned",
 "GetmanY1/wav2vec2-large-fi-150k-finetuned",
 "GetmanY1/wav2vec2-xlarge-fi-150k-finetuned"
]

gradio_app = gr.Interface(
    fn=transcribe,
    inputs=[
        gr.Audio(sources=["upload","microphone"]),
        gr.Dropdown(
            label="Model",
            value="GetmanY1/wav2vec2-large-fi-150k-finetuned",
            choices=model_ids_list
        ),
        gr.Checkbox(
            label="Add Punctuation and Capitalization*",
            value=True,
            info="*An external model (1-800-BAD-CODE/xlm-roberta_punctuation_fullstop_truecase) is used to restore punctuation and capitalization"
        )
        ],
    outputs="text",
    title="Finnish Automatic Speech Recognition",
    description ="Choose a model from the list. Select the Base model for the fastest inference and the XLarge one for the most accurate results."
)

if __name__ == "__main__":
    gradio_app.launch()