import streamlit as st
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM

def initialize_translator(model_name):
    return pipeline("translation", model=model_name)

model_name = "Helsinki-NLP/opus-mt-en-ru"
translator = initialize_translator(model_name)

def translate_text(text):
    if text:
        result = translator(text)
        return result[0]['translation_text']
    return ""

st.title("Text Translation App")

st.sidebar.header("Settings")
language_pair = st.sidebar.selectbox(
    "Choose language pair:",
    [
        "English to Russian (Helsinki-NLP/opus-mt-en-ru)",
        "Russian to English (Helsinki-NLP/opus-mt-ru-en)"
    ]
)

if "Russian to English" in language_pair:
    model_name = "Helsinki-NLP/opus-mt-ru-en"
else:
    model_name = "Helsinki-NLP/opus-mt-en-ru"

translator = initialize_translator(model_name)

st.subheader("Enter text to translate:")

user_input = st.text_area("Your text here (e.g., 'The weather is nice today.'):", height=200)

if st.button("Translate"):
    translation = translate_text(user_input)
    st.subheader("Translated Text:")
    st.write(translation)
else:
    st.info("Enter text and click 'Translate' to see the result.")


if __name__ == '__main__':
    import streamlit.web.cli as stcli
    stcli.main()