Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 import gradio as gr
-from transformers import MarianMTModel, MarianTokenizer, AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import MarianMTModel, MarianTokenizer, M2M100ForConditionalGeneration, M2M100Tokenizer
 
 # Отключение HF_HUB_ENABLE_HF_TRANSFER
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "false"
@@ -10,10 +10,10 @@ model_name_ka_ru = 'Helsinki-NLP/opus-mt-ka-ru'
 tokenizer_ka_ru = MarianTokenizer.from_pretrained(model_name_ka_ru, token=True)
 model_ka_ru = MarianMTModel.from_pretrained(model_name_ka_ru, token=True)
 
-# Модель для перевода с русского на грузинский (NLLB)
-model_name_ru_ka = "facebook/nllb-200-distilled-600M"
-tokenizer_ru_ka = AutoTokenizer.from_pretrained(model_name_ru_ka, token=True)
-model_ru_ka = AutoModelForSeq2SeqLM.from_pretrained(model_name_ru_ka, token=True)
+# Модель для перевода с русского на грузинский (M2M100)
+model_name_ru_ka = "facebook/m2m100_1.2B"
+tokenizer_ru_ka = M2M100Tokenizer.from_pretrained(model_name_ru_ka, token=True)
+model_ru_ka = M2M100ForConditionalGeneration.from_pretrained(model_name_ru_ka, token=True)
 
 def translate(text, direction):
     if direction == "Грузинский -> Русский":
@@ -22,17 +22,17 @@ def translate(text, direction):
         translated = model_ka_ru.generate(**inputs)
         translated_text = tokenizer_ka_ru.decode(translated[0], skip_special_tokens=True)
     else:
-        # Перевод с русского на грузинский (NLLB)
-        tokenizer_ru_ka.src_lang = "rus_Cyrl"  # Установка исходного языка на русский
-        target_lang_token_id = tokenizer_ru_ka.convert_tokens_to_ids(">>kau_KZ<<")  # Целевой язык: грузинский
+        # Перевод с русского на грузинский (M2M100)
+        tokenizer_ru_ka.src_lang = "ru"  # Установка исходного языка на русский
+        target_lang = "ka"  # Целевой язык: грузинский
 
         try:
             inputs = tokenizer_ru_ka(text, return_tensors="pt", padding=True, truncation=True)
             translated_tokens = model_ru_ka.generate(
                 **inputs,
-                forced_bos_token_id=target_lang_token_id,  # Правильный идентификатор грузинского языка
-                max_length=100,  # Уменьшение длины для ускорения
-                num_beams=2,  # Уменьшение лучей для увеличения скорости
+                forced_bos_token_id=tokenizer_ru_ka.get_lang_id(target_lang),  # Установка целевого языка
+                max_length=200,
+                num_beams=4,  # Оптимальное количество лучей для M2M100
                 early_stopping=True
             )
             translated_text = tokenizer_ru_ka.decode(translated_tokens[0], skip_special_tokens=True)
@@ -50,8 +50,8 @@ iface = gr.Interface(
         gr.Radio(["Грузинский -> Русский", "Русский -> Грузинский"], label="Выберите направление")
     ],
     outputs="text",
-    title="Русско-Грузинский Переводчик (Helsinki-NLP и NLLB, оптимизирован)",
-    description="Перевод с использованием модели Helsinki-NLP и NLLB"
+    title="Русско-Грузинский Переводчик (Helsinki-NLP и M2M100, оптимизированный)",
+    description="Перевод с использованием проверенных моделей Helsinki-NLP и M2M100"
 )
 
 iface.launch()
@@ -80,5 +80,6 @@ iface.launch()
 
 
 
+