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

# Function to translate text
@st.cache
def translation(txt):
    pipe = pipeline("translation_en_to_zh", model="DDDSSS/translation_en-zh")
    rst = pipe(txt, max_length=40)
    return rst

# Main function
def main():
    uploaded_file = st.file_uploader("Select an audio file")
    
    if uploaded_file is not None:
        bytes_data = uploaded_file.getvalue()
        with open(uploaded_file.name, "wb") as file:
            file.write(bytes_data)
        
        st.image(uploaded_file, caption="Uploaded Audio", use_column_width=True)
        
        # Stage 1: Audio to Text
        st.text('Processing audio2txt...')
        txt = audio2txt(uploaded_file.name)
        st.write(txt)
        
        # Stage 2: Text to Translation
        st.text('Generating a translation...')
        txt2 = translation(txt)
        st.write(txt2)

if __name__ == "__main__":
    main()