import streamlit as st
from transformers import MarianMTModel, MarianTokenizer

# Function to load model and tokenizer
@st.cache_resource
def load_model_and_tokenizer(model_name):
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    model = MarianMTModel.from_pretrained(model_name)
    return tokenizer, model

# Function to perform translation
def translate_text(text, tokenizer, model):
    tokenized_text = tokenizer.prepare_seq2seq_batch([text], return_tensors="pt", padding=True)
    translated_tokens = model.generate(**tokenized_text)
    translated_text = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
    return translated_text

# Available language pairs (from Helsinki-NLP models)
language_pairs = {
    "English to French": "Helsinki-NLP/opus-mt-en-fr",
    "French to English": "Helsinki-NLP/opus-mt-fr-en",
    "English to German": "Helsinki-NLP/opus-mt-en-de",
    "German to English": "Helsinki-NLP/opus-mt-de-en",
    "English to Spanish": "Helsinki-NLP/opus-mt-en-es",
    "Spanish to English": "Helsinki-NLP/opus-mt-es-en",
    # Add more pairs as needed
}

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

# User selects language pair
language_pair = st.selectbox("Select Language Pair:", list(language_pairs.keys()))
model_name = language_pairs[language_pair]

# Load model and tokenizer
with st.spinner("Loading translation model..."):
    tokenizer, model = load_model_and_tokenizer(model_name)

# Input text
input_text = st.text_area("Enter text to translate:")

if st.button("Translate"):
    if input_text.strip():
        with st.spinner("Translating..."):
            translated_text = translate_text(input_text, tokenizer, model)
        st.success("Translation complete!")
        st.text_area("Translated Text:", translated_text, height=200)
    else:
        st.warning("Please enter text to translate.")