import streamlit as st
from transformers import MarianMTModel, MarianTokenizer

# Title of the app
st.title("🌍 Fun Translator App")
st.write("Translate text between languages instantly! Powered by Hugging Face Transformers 🤗")

# Language model mapping
LANGUAGE_MODELS = {
    "English to French": "Helsinki-NLP/opus-mt-en-fr",
    "English to German": "Helsinki-NLP/opus-mt-en-de",
    "English to Spanish": "Helsinki-NLP/opus-mt-en-es",
    "French to English": "Helsinki-NLP/opus-mt-fr-en",
    "German to English": "Helsinki-NLP/opus-mt-de-en",
    "Spanish to English": "Helsinki-NLP/opus-mt-es-en"
}

# Sidebar for language selection
selected_translation = st.sidebar.selectbox("Select Translation Pair:", list(LANGUAGE_MODELS.keys()))
model_name = LANGUAGE_MODELS[selected_translation]

# Input text area
input_text = st.text_area("Enter the text you want to translate:")

# Load model and tokenizer
@st.cache_resource
def load_model_and_tokenizer(model_name):
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    model = MarianMTModel.from_pretrained(model_name)
    return tokenizer, model

if st.button("Translate"):
    if input_text.strip():
        try:
            tokenizer, model = load_model_and_tokenizer(model_name)
            # Tokenize and translate
            inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512)
            outputs = model.generate(**inputs)
            translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
            st.success("Translated Text:")
            st.write(translated_text)
        except Exception as e:
            st.error(f"An error occurred: {e}")
    else:
        st.warning("Please enter some text to translate.")