import streamlit as st
from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer

# App Title and Description
st.title("🌐 Universal Language Translator App")
st.write("""
Translate text from any language to any other language using the open-source M2M100 multilingual model.  
This app supports over 100 languages and provides a seamless translation experience.
""")

# Instructions
st.sidebar.header("📝 Instructions")
st.sidebar.write("""
1. Enter the text you want to translate in the input box.
2. Select the source language of the input text.
3. Choose the target language for translation.
4. Click the "Translate" button to see the translated text.
""")

# Supported Languages
languages = {
    "English": "en",
    "Spanish": "es",
    "French": "fr",
    "German": "de",
    "Chinese": "zh",
    "Hindi": "hi",
    "Arabic": "ar",
    "Russian": "ru",
    "Italian": "it",
    "Portuguese": "pt",
    "Japanese": "ja",
    "Korean": "ko",
    "Dutch": "nl",
    "Bengali": "bn",
    "Turkish": "tr",
    "Urdu": "ur",
    "Greek": "el",
    "Polish": "pl",
    "Thai": "th",
    "Vietnamese": "vi",
    "Hebrew": "he",
    "Swahili": "sw",
    "Tamil": "ta",
    "Telugu": "te",
    "Punjabi": "pa",
    "Malayalam": "ml",
}

# UI for Input and Language Selection
source_text = st.text_area("Enter text to translate:", height=150, placeholder="Type here...")
source_language = st.selectbox("Source Language:", options=list(languages.keys()), index=0)
target_language = st.selectbox("Target Language:", options=list(languages.keys()), index=1)

if st.button("Translate"):
    if source_text.strip() == "":
        st.error("Please enter some text to translate.")
    else:
        try:
            # Load M2M100 model and tokenizer
            model_name = "facebook/m2m100_418M"
            tokenizer = M2M100Tokenizer.from_pretrained(model_name)
            model = M2M100ForConditionalGeneration.from_pretrained(model_name)
            
            # Set source and target language
            tokenizer.src_lang = languages[source_language]
            encoded_text = tokenizer(source_text, return_tensors="pt")
            generated_tokens = model.generate(**encoded_text, forced_bos_token_id=tokenizer.get_lang_id(languages[target_language]))
            
            # Decode the translated text
            translated_text = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
            
            # Display Translated Text
            st.subheader("🔄 Translated Text:")
            st.text_area("Translation Output:", value=translated_text, height=150)
        except Exception as e:
            st.error("Translation failed. Ensure the text and languages are valid.")
            st.error(str(e))

# Footer
st.write("---")
st.markdown("💡 **Developed by Abdullah**")