import streamlit as st
from transformers import pipeline, AutoTokenizer

base="Helsinki-NLP/opus-mt-en-zh"
model="edwinlaw/opus-mt-cantonese-v2"
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:")
st.markdown(
    """
    This translation engine is a fine-tuned version of the [Helsinki-NLP/opus-mt-en-zh](https://huggingface.co/Helsinki-NLP/opus-mt-en-zh) model, initially designed for translating English to Mandarin Chinese.
    
    This [fine-tuned model](https://huggingface.co/edwinlaw/opus-mt-cantonese-v2) has been trained using 6000+ Cantonese sentences along with their English translations.
    
    Translations generated are intended for experimental use only. Correctness is not guaranteed. Proceed at your own risk.
    """)

st.subheader('Enter a simple English sentence here:')
prompt = st.text_input('')

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)