Add application file

app.py

@@ -1,31 +1,43 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
 
-# Загружаем модели
-pipe_kz_to_ru = pipeline(model="deepvk/kazRush-kk-ru")  # Казахский → Русский
-pipe_ru_to_kz = pipeline(model="deepvk/kazRush-ru-kk")  # Русский → Казахский
+# Загружаем модель и токенизатор
+model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
+tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
 
-# Название и описание приложения
-st.title("Переводчик: Казахский ↔ Русский")
-st.write("Выберите направление перевода и введите текст.")
+# Функция перевода
+def translate_text(text, src_lang, tgt_lang):
+    tokenizer.src_lang = src_lang
+    encoded = tokenizer(text, return_tensors="pt")
+    generated_tokens = model.generate(
+        **encoded,
+        forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang]
+    )
+    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
 
-# Выпадающий список для выбора направления перевода
-language_direction = st.selectbox(
-    "Выберите направление перевода:",
-    ["Казахский → Русский", "Русский → Казахский"]
-)
+# Название и описание
+st.title("Переводчик MBart")
+st.write("Перевод между русским и английским языками.")
 
-# Ввод текста пользователем
-input_text = st.text_area("Введите текст:")
+# Ввод текста и выбор языков
+input_text = st.text_area("Введите текст для перевода:")
 
-# Обработка перевода при нажатии кнопки
+language_options = {
+    "Русский": "ru_RU",
+    "Английский": "en_XX"
+}
+
+# Пользователь выбирает исходный и целевой язык
+src_lang = st.selectbox("Выберите исходный язык:", list(language_options.keys()))
+tgt_lang = st.selectbox("Выберите целевой язык:", list(language_options.keys()))
+
+# Перевод при нажатии кнопки
 if st.button("Перевести"):
-    if input_text.strip():
-        if language_direction == "Казахский → Русский":
-            result = pipe_kz_to_ru(input_text)
-        else:  # Русский → Казахский
-            result = pipe_ru_to_kz(input_text)
-        
-        st.write("Перевод:", result[0]['translation_text'])
+    if tgt_lang == src_lang:
+        st.write("Выберите только другой язык")
+    elif input_text.strip():
+        translation = translate_text(input_text, language_options[src_lang], language_options[tgt_lang])
+        st.write(f"Перевод с {src_lang} на {tgt_lang}:")
+        st.write(translation[0])
     else:
         st.write("Пожалуйста, введите текст для перевода.")