import os
import gradio as gr
from transformers import MarianMTModel, MarianTokenizer, M2M100ForConditionalGeneration, M2M100Tokenizer

# Отключение HF_HUB_ENABLE_HF_TRANSFER
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "false"

# Модель для перевода с грузинского на русский (Helsinki-NLP)
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)

# Модель для перевода с русского на грузинский (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 == "Грузинский -> Русский":
        # Перевод с грузинского на русский
        inputs = tokenizer_ka_ru(text, return_tensors="pt", padding=True, truncation=True)
        translated = model_ka_ru.generate(**inputs)
        translated_text = tokenizer_ka_ru.decode(translated[0], skip_special_tokens=True)
    else:
        # Перевод с русского на грузинский (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=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)

        except Exception as e:
            translated_text = f"Ошибка: {str(e)}"

    return translated_text

# Интерфейс Gradio
iface = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(label="Введите текст"),
        gr.Radio(["Грузинский -> Русский", "Русский -> Грузинский"], label="Выберите направление")
    ],
    outputs="text",
    title="Русско-Грузинский Переводчик (Helsinki-NLP и M2M100, оптимизированный)",
    description="Перевод с использованием проверенных моделей Helsinki-NLP и M2M100"
)

iface.launch()