import streamlit as st
from transformers import pipeline

# Define model mapping for each target language
model_mapping = {
    "French": "Helsinki-NLP/opus-mt-en-fr",
    "Spanish": "Helsinki-NLP/opus-mt-en-es",
    "German": "Helsinki-NLP/opus-mt-en-de",
    "Italian": "Helsinki-NLP/opus-mt-en-it",
    "Urdu": "Helsinki-NLP/opus-mt-en-ur"  # Added Urdu translation
}

# Streamlit app
st.title("🌎 TRANSLINGO 🔡")

st.markdown(
    """
    <style>
    .main {
        background: linear-gradient(to right, #ff7e5f, #feb47b);
        height: 100vh;
        padding: 2rem;
    }
    .block-container {
        padding: 2rem;
    }
    .title {
        background-color: #d4edda;  /* Light green color */
        padding: 1rem;
        border-radius: 8px;
        text-align: center;
    }
    </style>
    """,
    unsafe_allow_html=True
)


col1, col2, col3 = st.columns([1, 3, 1])

# Center the content in the middle column
with col2:
    st.header("Translate Text")
    
    # Input text
    text = st.text_area("Enter text in English:")
    
    # Target language selection
    selected_language = st.selectbox("Select target language:", list(model_mapping.keys()))
    
    # Translation button
    if st.button("Translate"):
        if text:
            # Get the appropriate model based on selected language
            model_name = model_mapping[selected_language]
            translator = pipeline("translation", model=model_name)
            
            # Perform translation
            translation = translator(text, max_length=400)[0]['translation_text']
            
            # Display translation
            st.write(f"*Translation in {selected_language}:*")
            st.write(translation)
        else:
            st.write("Please enter some text to translate.")