import streamlit as st
from transformers import MarianMTModel, MarianTokenizer

# Initialize the model and tokenizer
def load_model_and_tokenizer(target_lang):
    model_name = f'Helsinki-NLP/opus-mt-en-{target_lang}'
    model = MarianMTModel.from_pretrained(model_name)
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    return model, tokenizer

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

# Streamlit app
st.title("Language Translator")

# Input text
input_text = st.text_area("Enter text to translate", "Hello, how are you?")

# Language selection
target_language = st.selectbox(
    "Select target language",
    ["fr", "de", "es", "it", "pt", "ru", "zh", "ja", "ar"]
)

# Load model and tokenizer based on selected language
if target_language:
    model, tokenizer = load_model_and_tokenizer(target_language)
    
    if st.button("Translate"):
        if input_text:
            translated_text = translate_text(input_text, model, tokenizer)
            st.subheader("Translated text")
            st.write(translated_text)
        else:
            st.error("Please enter some text to translate.")

streamlit run app.py