import streamlit as st
from transformers import pipeline, AutoTokenizer

base="Helsinki-NLP/opus-mt-en-zh"
model="edwinlaw/opus-mt-cantonese-v1"
tokenizer = AutoTokenizer.from_pretrained(base)

def translate(text, src_lang, tgt_lang):
    translator = pipeline(
        "translation",
        model=model,
        tokenizer=tokenizer,
        src_lang=src_lang,
        tgt_lang=tgt_lang,
    )
    translated_text = translator(text)
    return translated_text

st.title("Translate English into Cantonese:")

prompt = st.text_input('English sentence here')

if prompt:
    translation = translate(prompt, 'en', 'yue')
    translated_txt = translation[0]['translation_text']

    st.write(translated_txt)

    with st.expander('Chat History'):
        st.info(translated_txt)