import streamlit as st
from transformers import MarianMTModel, MarianTokenizer

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

# Language selection options (you can expand this list as needed)
language_options = {
    "English to French": ("en", "fr"),
    "French to English": ("fr", "en"),
    "English to Spanish": ("en", "es"),
    "Spanish to English": ("es", "en"),
    "English to German": ("en", "de"),
    "German to English": ("de", "en"),
    "English to Italian": ("en", "it"),
    "Italian to English": ("it", "en"),
    # Add more language pairs as needed
}

# Streamlit UI setup
st.title("Hugging Face Translation App")
st.write("This app allows you to translate text between different languages using Hugging Face's translation models.")

# Select the language pair
language_pair = st.selectbox("Select language pair", list(language_options.keys()))
src_lang, tgt_lang = language_options[language_pair]

# Load the corresponding model and tokenizer
model_name = f"Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}"
try:
    model, tokenizer = load_model_and_tokenizer(model_name)
except Exception as e:
    st.error(f"Error loading model {model_name}: {e}")
    st.stop()  # Stop execution if the model loading fails

# Input text to translate
text = st.text_area("Enter text to translate")

# Translate button
if st.button("Translate"):
    if text.strip():
        # Prepare the input for the model
        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
        
        # Generate translation
        translation = model.generate(**inputs)
        
        # Decode the output
        translated_text = tokenizer.decode(translation[0], skip_special_tokens=True)
        
        # Display translation
        st.write(f"**Original Text**: {text}")
        st.write(f"**Translated Text**: {translated_text}")
    else:
        st.warning("Please enter some text to translate.")