import streamlit as st
from transformers import MarianMTModel, MarianTokenizer

# Function to perform translation
def translate(text, src_lang_code, tgt_lang_code):
    model_name = f'Helsinki-NLP/opus-mt-{src_lang_code}-{tgt_lang_code}'
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    model = MarianMTModel.from_pretrained(model_name)

    # Tokenize and translate
    inputs = tokenizer.encode(text, return_tensors="pt", max_length=512, truncation=True)
    translated_tokens = model.generate(inputs, max_length=512, num_beams=4, early_stopping=True)
    translated_text = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)

    return translated_text

# Language mapping
language_map = {
    "English": "en",
    "French": "fr",
    "German": "de",
    "Spanish": "es",
    "Italian": "it",
    "Dutch": "nl",
    "Chinese": "zh",
    "Hindi": "hi"
}

# Streamlit app layout
st.title("Language Translator App")

# Input and output language options
languages_full = list(language_map.keys())

# Language selection
src_lang_full = st.selectbox("Select input language:", languages_full, index=0)
tgt_lang_full = st.selectbox("Select output language:", languages_full, index=1)

# Map full names to language codes
src_lang_code = language_map[src_lang_full]
tgt_lang_code = language_map[tgt_lang_full]

# Text input
text_to_translate = st.text_area("Enter text to translate:")

# Translate button
if st.button("Translate"):
    if src_lang_code == tgt_lang_code:
        st.warning("Input and output languages must be different.")
    elif text_to_translate.strip() == "":
        st.warning("Please enter text to translate.")
    else:
        with st.spinner("Translating..."):
            try:
                translation = translate(text_to_translate, src_lang_code, tgt_lang_code)
                st.success("Translation completed!")
                st.text_area("Translated text:", value=translation, height=200)
            except Exception as e:
                st.error(f"An error occurred: {e}")