from speechbox import PunctuationRestorer
import librosa
import subprocess
import gradio as gr

restorer = PunctuationRestorer.from_pretrained("openai/whisper-tiny.en")


def convert_to_wav(path):
    if path[-3:] != 'wav':
        new_path = '.'.join(path.split('.')[:-1]) + '.wav'
    try:
        subprocess.call(['ffmpeg', '-i', path, new_path, '-y'])
    except:  # noqa: E722
        return path, 'Error: Could not convert file to .wav'
    path = new_path
    return path, None


def restore(audio, original_transcript):
    path, error = convert_to_wav(audio)
    print(error)
    data, samplerate = librosa.load(path, sr=16_000)

    text, log_probs = restorer(data, original_transcript, samplerate, num_beams=1)

    return text, log_probs


gr.Interface(
    title='Punctuation Restorer',
    fn=restore,
    inputs=[
        gr.inputs.Audio(source="upload", type="filepath"),
        gr.inputs.Textbox(default="", label="normalized text")
    ],
    outputs=[
        gr.outputs.Textbox(label='Restored text'),
        gr.Number(label='Log probability')
    ],
    examples=[
        ["./common_voice_en_18301577.mp3", "do not cross the yellow light"],
        ["./sample1.flac", "going along slushy country roads and speaking to damp audiences in draughty school rooms day after day for a fortnight he'll have to put in an appearance at some place of worship on sunday morning and he can come to us immediately afterwards"],
        ["./sample2.flac", "before he had time to answer a much encumbered vera burst into the room with the question i say can i leave these here these were a small black pig and a lusty specimen of black red game cock"],
    ]
  ).launch()