app.py

import streamlit as st
from transformers import MarianMTModel, MarianTokenizer

# Define available languages with MarianMT models
LANGUAGES = {
    'Spanish': 'es',
    'French': 'fr',
    'German': 'de',
    'Chinese': 'zh',
    'Hindi': 'hi',
    'Arabic': 'ar',
    'Japanese': 'ja',
    'Russian': 'ru',
    'Italian': 'it',
    'Portuguese': 'pt',
    # Add more languages if needed
}

# Function to load the model based on the selected language
@st.cache_resource
def load_model(src_lang='en', tgt_lang='es'):
    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

# Function to translate text
def translate_text(model, tokenizer, text):
    inputs = tokenizer.encode(text, return_tensors='pt', truncation=True, padding=True)
    translated = model.generate(inputs, max_length=512, num_beams=5, early_stopping=True)
    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
    return translated_text

# Streamlit app
st.title("Language Translator")
st.write("Translate English text to any language using Hugging Face models.")

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

# Language selection
language = st.selectbox("Choose target language", list(LANGUAGES.keys()))

if st.button("Translate"):
    if text:
        # Load model and tokenizer based on selected language
        tgt_lang = LANGUAGES[language]
        model, tokenizer = load_model('en', tgt_lang)
        
        # Perform translation
        translated_text = translate_text(model, tokenizer, text)
        
        # Display the translation
        st.write(f"**Translated text ({language}):**")
        st.write(translated_text)
    else:
        st.write("Please enter text to translate.")