import streamlit as st
import time
from transformers import pipeline

st.title("Enhanced Multilingual Translator Chatbot")

# Choose the translation models from Hugging Face
translation_models = {
    "English to German": "Helsinki-NLP/opus-mt-en-de",
    "German to English": "Helsinki-NLP/opus-mt-de-en",
    "English to French": "Helsinki-NLP/opus-mt-en-fr",
    "French to English": "Helsinki-NLP/opus-mt-fr-en",
    "English to Urdu": "Helsinki-NLP/opus-mt-en-ur",
    "Urdu to English": "Helsinki-NLP/opus-mt-ur-en",
    "English to Spanish": "Helsinki-NLP/opus-mt-en-es",
    "Spanish to English": "Helsinki-NLP/opus-mt-es-en",
    "English to Chinese": "Helsinki-NLP/opus-mt-en-zh",
    "Chinese to English": "Helsinki-NLP/opus-mt-zh-en",
    # Add more language pairs as needed
}

selected_translation = st.selectbox("Select translation model", list(translation_models.keys()))

# Load the translation pipeline
translator = pipeline(task="translation", model=translation_models[selected_translation])

# User input for translation
user_input = st.text_area("Enter text for translation:", "")

# Display loading indicator
if st.button("Translate"):
    with st.spinner("Translating..."):
        # Simulate translation delay for demonstration
        time.sleep(2)
        if user_input:
            # Perform translation
            translated_text = translator(user_input, max_length=500)[0]['translation_text']
            st.success(f"Translated Text: {translated_text}")
        else:
            st.warning("Please enter text for translation.")

# Clear button to reset input and result
if st.button("Clear"):
    user_input = ""
    st.success("Input cleared.")
    st.empty()  # Clear previous results if any

st.markdown("---")

st.subheader("About")
st.write(
    "This is an enhanced Multilingual Translator chatbot that uses the Hugging Face Transformers library."
)
st.write(
    "Select a translation model from the dropdown, enter text, and click 'Translate' to see the translation."
)