# pip install streamlit
# pip install audio-recorder-streamlit

import streamlit as st
from st_audiorec import st_audiorec
import nemo.collections.asr as nemo_asr
from pydub import AudioSegment
import io
import os
import uuid 

@st.cache_resource
def get_model():
    try:
        os.makedirs("audio_cache")
    except: 
        pass
    en_asr_model = nemo_asr.models.EncDecCTCModelBPE.restore_from("/home/akshat/v1/en_am_model/Conformer-CTC-BPE-Large.nemo")

    return en_asr_model

en_asr_model = get_model()

st.title("💬 Vocalize: Empower Your Voice ")

"""
Hi record the audio, and get the transcription in real time!
Note: Works best for smaller audios
"""

wav_audio_data = st_audiorec()

if wav_audio_data:
    audio_location = "audio_cache/" + str(uuid.uuid4()) + ".wav"
    audio_file = io.BytesIO(wav_audio_data)
    audio = AudioSegment.from_file(audio_file)
    audio = audio.set_sample_width(2)
    audio = audio.set_channels(1)
    audio = audio.set_frame_rate(16000)
    audio.export(audio_location, format="wav")

    text = en_asr_model.transcribe([audio_location], logprobs=False)[0]
    print(text)
    st.write(text)