import streamlit as st
from transformers import MarianMTModel, MarianTokenizer
import torch

# Title and Description
st.set_page_config(page_title='Language Translation App', layout='centered')
st.title('Language Translation App')
st.write('Translate text between multiple languages using an open-source model.')

# Model Selection
@st.cache_resource
def load_model(src_lang, tgt_lang):
    if not torch.cuda.is_available():
        st.warning('CUDA is not available. The model will run on CPU, which may be slower.')
    model_name = f'Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}'
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    model = MarianMTModel.from_pretrained(model_name)
    return tokenizer, model

# Sidebar for Language Selection
st.sidebar.header('Language Selection')
src_lang = st.sidebar.selectbox('Select Source Language', ['en', 'fr', 'es', 'de', 'it', 'pt', 'ru', 'zh', 'ja', 'ar'], index=0)
tgt_lang = st.sidebar.selectbox('Select Target Language', ['en', 'fr', 'es', 'de', 'it', 'pt', 'ru', 'zh', 'ja', 'ar'], index=1)

# Main Input Area
st.write('### Enter text to translate:')
text = st.text_area('Text Input', placeholder='Type your text here...')

# Translation Button
if st.button('Translate'):
    if text.strip() == '':
        st.warning('Please enter text to translate.')
    else:
        try:
            tokenizer, model = load_model(src_lang, tgt_lang)
            inputs = tokenizer.encode(text, return_tensors='pt')
            translated = model.generate(inputs, max_length=1000, num_beams=5, early_stopping=True)
            result = tokenizer.decode(translated[0], skip_special_tokens=True)
            st.success('### Translated Text:')
            st.write(result)
        except Exception as e:
            st.error(f'Error: {e}')

# Footer
st.sidebar.write('Powered by Helsinki-NLP/opus-mt')
st.sidebar.write('Built with Streamlit')

# Instructions for Running
st.write('**To run the app:**')
st.code('streamlit run app.py')