import os
import streamlit as st
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, TranslationPipeline

print("Loading the model...")

hf_token = os.getenv("HF_AUTH_TOKEN")
if not hf_token:
    raise ValueError("Hugging Face token not found. Please set the HF_AUTH_TOKEN environment variable.")

# Title and Description
st.title("Translation Web App")
st.write("""
### Powered by Hugging Face and Streamlit
This app uses a pre-trained NLP model from Hugging Face to translate text between languages.  
Enter text in the source language, select source and target languages, and see the translation!
""")

# Initialize Hugging Face Translation Pipeline
@st.cache_resource
def load_translation_pipeline():
    print("Loading translation model...")
    model = AutoModelForSeq2SeqLM.from_pretrained(
        'issai/tilmash',
        use_auth_token=hf_token
    )
    tokenizer = AutoTokenizer.from_pretrained(
        "issai/tilmash",
        use_auth_token=hf_token
    )
    return TranslationPipeline(model=model, tokenizer=tokenizer, max_length=1000)

tilmash = load_translation_pipeline()

languages = {
    "Kazakh (Cyrillic)": "kaz_Cyrl",
    "Russian (Cyrillic)": "rus_Cyrl",
    "English (Latin)": "eng_Latn",
    "Turkish (Latin)": "tur_Latn"
}

src_lang = st.selectbox("Select source language:", options=list(languages.keys()), index=0)
tgt_lang = st.selectbox("Select target language:", options=list(languages.keys()), index=2)

user_input = st.text_area("Enter text to translate:", "")

if st.button("Translate Text"):
    if user_input.strip():
        result = tilmash(user_input, src_lang=languages[src_lang], tgt_lang=languages[tgt_lang])
        translation = result[0]['translation_text']

        st.subheader("Translation Result")
        st.write(f"**Translated Text:** {translation}")
    else:
        st.warning("Please enter some text to translate!")

# Sidebar with About Information
st.sidebar.title("About")
st.sidebar.info("""
This app demonstrates the use of Hugging Face's NLP models with Streamlit.  
It uses the `issai/tilmash` model for translation between languages such as Kazakh, Russian, English, and Turkish.
""")

print('After translation operation')