import streamlit as st
from transformers import MarianMTModel, MarianTokenizer

st.set_page_config(page_title="Language Translator", layout="wide")

# Function to load model and tokenizer based on selected languages
def load_model_and_tokenizer(src_lang, tgt_lang):
    model_name = f"Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}"
    model = MarianMTModel.from_pretrained(model_name)
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    return model, tokenizer

# Function to translate text
def translate_text(text, model, tokenizer):
    inputs = tokenizer(text, return_tensors="pt", padding=True)
    translated = model.generate(**inputs)
    return tokenizer.decode(translated[0], skip_special_tokens=True)

# Streamlit app layout
st.title("Language Translator App")
st.write("Translate text between multiple languages using open-source models.")

languages = {
    "English": "en",
    "French": "fr",
    "Spanish": "es",
    "German": "de",
    "Italian": "it",
    "Russian": "ru",
    "Chinese": "zh"
}

source_language = st.selectbox("Select Input Language:", options=list(languages.keys()))
target_language = st.selectbox("Select Output Language:", options=list(languages.keys()))

if source_language == target_language:
    st.warning("Input and output languages must be different.")
else:
    src_lang_code = languages[source_language]
    tgt_lang_code = languages[target_language]

    with st.spinner("Loading translation model..."):
        model, tokenizer = load_model_and_tokenizer(src_lang_code, tgt_lang_code)

    col1, col2 = st.columns(2)

    with col1:
        text_to_translate = st.text_area("Enter text to translate:", height=300)

    with col2:
        if st.button("Translate"):
            if text_to_translate.strip():
                with st.spinner("Translating..."):
                    translated_text = translate_text(text_to_translate, model, tokenizer)
                    st.success("Translation Completed!")
                    st.text_area("Translated Text:", value=translated_text, height=300)
            else:
                st.error("Please enter text to translate.")