import streamlit as st
from transformers import MarianMTModel, MarianTokenizer

# Supported languages
supported_languages = {
    'English to Urdu': 'Helsinki-NLP/opus-mt-en-ur',
    'English to French': 'Helsinki-NLP/opus-mt-en-fr',
    'English to Spanish': 'Helsinki-NLP/opus-mt-en-es',
    'French to English': 'Helsinki-NLP/opus-mt-fr-en',
    'Spanish to English': 'Helsinki-NLP/opus-mt-es-en',
    'Urdu to English': 'Helsinki-NLP/opus-mt-ur-en'
}

# Streamlit App Title
st.title("Multi-Language Translator")

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

# Language selection dropdown
language_pair = st.selectbox("Choose the language pair:", list(supported_languages.keys()))

# Load the model and tokenizer for the selected language pair
if input_text:
    model_name = supported_languages[language_pair]
    model = MarianMTModel.from_pretrained(model_name)
    tokenizer = MarianTokenizer.from_pretrained(model_name)

    # Tokenize and translate
    translated = model.generate(**tokenizer(input_text, return_tensors="pt", padding=True))
    translated_text = [tokenizer.decode(t, skip_special_tokens=True) for t in translated]

    # Display the translated text
    st.subheader("Translated Text:")
    st.write(translated_text[0])

# Footer
st.markdown("Powered by Helsinki-NLP/Opus-MT and Hugging Face")