import streamlit as st
from transformers import pipeline

# Function to initialize the translation pipeline
@st.cache_resource
def load_translation_pipeline(source_lang, target_lang):
    model_name = f"Helsinki-NLP/opus-mt-{source_lang}-{target_lang}"
    return pipeline("translation", model=model_name)

# Function to perform translation
def translate_text(text, source_lang, target_lang):
    translator = load_translation_pipeline(source_lang, target_lang)
    translated = translator(text, max_length=512)
    return translated[0]['translation_text']

# Streamlit app
st.title("Language Translation App")
st.markdown("Translate text quickly between multiple languages using open-source models.")

# Input text area
text_to_translate = st.text_area("Enter text to translate:", placeholder="Type here...", height=150)

# Language selection options
languages = {
    "English": "en",
    "French": "fr",
    "German": "de",
    "Spanish": "es",
    "Italian": "it",
    "Chinese": "zh",
    "Hindi": "hi",
    "Urdu": "ur",
    "Persian": "fa"
}

# Select source and target languages
source_language = st.selectbox("Select source language:", list(languages.keys()))
target_language = st.selectbox("Select target language:", list(languages.keys()))

# Get language codes
source_lang_code = languages[source_language]
target_lang_code = languages[target_language]

# Translate button
if st.button("Translate"):
    if not text_to_translate.strip():
        st.warning("Please enter some text to translate.")
    elif source_lang_code == target_lang_code:
        st.warning("Source and target languages must be different.")
    else:
        try:
            translation = translate_text(text_to_translate, source_lang_code, target_lang_code)
            st.success("Translation completed!")
            st.text_area("Translated Text:", translation, height=150, disabled=True)
        except Exception as e:
            st.error(f"An error occurred during translation: {str(e)}")

# Footer
st.markdown("---")
st.markdown(
    "**Note:** This app uses Hugging Face's Helsinki-NLP models for translation. These models are open-source and ideal for various language translation tasks."
)