import streamlit as st
from transformers import MarianMTModel, MarianTokenizer

# Define supported language pairs (source-target)
LANGUAGES = {
    "en": "English",
    "fr": "French",
    "de": "German",
    "es": "Spanish",
    "it": "Italian",
    "pt": "Portuguese",
    "ru": "Russian",
    "zh": "Chinese",
}

def get_model_and_tokenizer(src_lang: str, tgt_lang: str):
    model_name = f"Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}"
    model = MarianMTModel.from_pretrained(model_name)
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    return model, tokenizer

def translate_text(text: str, source_lang: str, target_lang: str):
    model, tokenizer = get_model_and_tokenizer(source_lang, target_lang)
    translated = model.generate(**tokenizer(text, return_tensors="pt", padding=True))
    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
    return translated_text

# Streamlit app
st.title("Text Translation App")

st.write("Translate text from one language to another using a Hugging Face model.")

source_lang = st.selectbox("Select the source language", list(LANGUAGES.keys()), format_func=lambda x: LANGUAGES[x])
target_lang = st.selectbox("Select the target language", list(LANGUAGES.keys()), format_func=lambda x: LANGUAGES[x])

text = st.text_area("Enter the text you want to translate")

if st.button("Translate"):
    if source_lang == target_lang:
        st.warning("Source and target languages are the same. Please select different languages.")
    else:
        translated_text = translate_text(text, source_lang, target_lang)
        st.write("**Translated Text:**")
        st.write(translated_text)