import streamlit as st
from transformers import pipeline

st.set_page_config(page_title="Your English audio to Chinese text", page_icon="🦜")
st.header("Turn Your English Audio to Chinese text")

# Function to convert audio to text
@st.cache
def audio2txt(audioname):
    pipe = pipeline("automatic-speech-recognition", model="avery0/pipeline1model2")
    rst = pipe(audioname)
    return rst['text']

# Function to translate text
@st.cache
def translation(txt):
    pipe = pipeline("translation", model= "Helsinki-NLP/opus-mt-en-zh")
    rst = pipe(txt)
    return rst

# Main function
def main():
    uploaded_file = st.file_uploader("Select an audio file", type=["mp3", "wav","m4a"])
    
    if uploaded_file is not None:
        audio_data = uploaded_file.read()
        st.audio(audio_data, format='audio/mp3/m4a')
        
        # Stage 1: Audio to Text
        st.text('Processing audio2txt...')
        txt = audio2txt(audio_data)
        st.write(txt)
        
        # Stage 2: Text to Translation
        st.text('Generating a translation...')
        txt2 = translation(txt)
        st.write(txt2)

if __name__ == "__main__":
    main()