import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
import re


# --------------------------
# Очистка текста
# --------------------------
def clean_text(s):
    s = s.strip()
    s = re.sub(r"\s+", " ", s)
    s = re.sub(r"\s+([,.!?;:])", r"\1", s)
    return s


# --------------------------
# Загрузка модели
# --------------------------
def load_model(model_name):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
    return tokenizer, model


# --------------------------
# Перевод
# --------------------------
def translate_text(text, model_name):
    text = clean_text(text)

    tokenizer, model = load_model(model_name)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)

    enc = tokenizer([text], return_tensors="pt", padding=True, truncation=True)
    enc = {k: v.to(device) for k, v in enc.items()}

    out = model.generate(**enc, max_length=150, num_beams=4)
    translated = tokenizer.decode(out[0], skip_special_tokens=True)

    return translated


# --------------------------
# Gradio Interface
# --------------------------
def build_ui():
    title = "🌐 Automatic Text Translator"
    description = """
    Лёгкий переводчик на HuggingFace Transformers.  
    Выберите модель → введите текст → получите перевод.
    """

    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown(f"# {title}")
        gr.Markdown(description)

        with gr.Row():
            inp = gr.Textbox(label="Введите текст", placeholder="Hello, how are you?")
            out = gr.Textbox(label="Перевод")

        model_name = gr.Dropdown(
            label="Модель перевода",
            value="Helsinki-NLP/opus-mt-en-ru",
            choices=[
                "Helsinki-NLP/opus-mt-en-ru",
                "Helsinki-NLP/opus-mt-ru-en",
                "Helsinki-NLP/opus-mt-en-de",
                "Helsinki-NLP/opus-mt-en-fr",
                "Helsinki-NLP/opus-mt-en-kaz"
            ],
        )

        btn = gr.Button("Перевести")

        btn.click(translate_text, inputs=[inp, model_name], outputs=out)

    return demo


demo = build_ui()

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