import gradio as gr
from transformers import MarianMTModel, MarianTokenizer
import torch

# Загрузка модели и токенизатора
MODEL_NAME = "Helsinki-NLP/opus-mt-en-ru"
tokenizer = MarianTokenizer.from_pretrained(MODEL_NAME)
model = MarianMTModel.from_pretrained(MODEL_NAME)

def translate(text):
    if not text.strip():
        return "Пожалуйста, введите текст для перевода."
    if len(text) > 500:
        return "Слишком длинный текст. Пожалуйста, введите до 500 символов."
    try:
        batch = tokenizer([text], return_tensors="pt", padding=True)
        with torch.no_grad():
            gen = model.generate(**batch)
        translated = tokenizer.batch_decode(gen, skip_special_tokens=True)[0]
        return translated
    except Exception as e:
        return f"Ошибка при переводе: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown("# Переводчик с английского на русский")
    with gr.Row():
        input_text = gr.Textbox(label="Текст на английском", placeholder="Введите текст на английском языке")
    with gr.Row():
        translate_btn = gr.Button("Перевести")
    with gr.Row():
        output_text = gr.Textbox(label="Перевод на русский", interactive=False)

    translate_btn.click(translate, inputs=input_text, outputs=output_text)

if __name__ == "__main__":
    demo.launch()